From 43a54a7c3579b60cccfa8df34485e9c3ebf31093 Mon Sep 17 00:00:00 2001
From: Umang Mistry <umistry@aarete.com>
Date: Thu, 28 Mar 2024 13:48:46 -0500
Subject: [PATCH] Added stored proc to client dag

---
 airflow/dags/client_name_dag.py | 104 +++++++++++++++++++++++---------
 1 file changed, 76 insertions(+), 28 deletions(-)

diff --git a/airflow/dags/client_name_dag.py b/airflow/dags/client_name_dag.py
index 5f18dc2..52b5b28 100644
--- a/airflow/dags/client_name_dag.py
+++ b/airflow/dags/client_name_dag.py
@@ -4,8 +4,15 @@ from airflow import DAG
 from airflow.operators.python_operator import PythonOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from datetime import datetime, timedelta
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+from airflow.operators.empty import EmptyOperator
+import logging
+from airflow.exceptions import AirflowFailException
 
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 SNOWFLAKE_CONN_ID = "doczy_dev_snowflake"
 DAG_ID = "load_client_config"
 DATABASE="DOCZY_DEV"
@@ -13,39 +20,48 @@ DATABASE="DOCZY_DEV"
 
 TAGS=["dev","config_interface","dataload"]
 
-default_params = {"client_list_file_name": ""}
 
 
-
-def process_csv_in_s3(params):
-    # Instantiate S3Hook
-    file_name = params['client_list_file_name']
+def process_csv_in_s3(**kwargs):
     s3_hook = S3Hook(aws_conn_id="aws_default")
-
     bucket_name = "doczy-dev-infra-raw-data-ingestion"
-    key = 'client_names_openair/' + file_name
-    # Get the file from S3
-    file_content = s3_hook.read_key(key, bucket_name)
+    key_prefix = 'client_names_openair/'
+
+    # List objects within the specified bucket and key prefix
+    objects = s3_hook.list_keys(bucket_name=bucket_name, prefix=key_prefix)
     
-    # Convert string to DataFrame
-    from io import StringIO
-    df = pd.read_csv(StringIO(file_content))
+    # Sort the objects by their last modified date and select the latest
+    if objects:
+        latest_file_key = sorted(objects)[-1]  # Assuming file names include a timestamp or incrementing number
 
-    # Apply the generate_s3_path function
-    df['s3_path'] = df['customer_name'].apply(generate_s3_path)
+        # Get the latest file from S3
+        file_content = s3_hook.read_key(latest_file_key, bucket_name)
+    
+        # Convert string to DataFrame
+        from io import StringIO
+        df = pd.read_csv(StringIO(file_content))
 
-    # Convert DataFrame to CSV string
-    csv_buffer = StringIO()
-    df.to_csv(csv_buffer, index=False)
-    csv_content = csv_buffer.getvalue()
+        # Apply the generate_s3_path function
+        df['s3_path'] = df['customer_name'].apply(generate_s3_path)
 
-    # Replace the file in S3 with the processed content
-    s3_hook.load_string(
-        string_data=csv_content,
-        key=key,
-        bucket_name=bucket_name,
-        replace=True
-    )
+        # Convert DataFrame to CSV string
+        csv_buffer = StringIO()
+        df.to_csv(csv_buffer, index=False)
+        csv_content = csv_buffer.getvalue()
+
+        # Replace the file in S3 with the processed content
+        s3_hook.load_string(
+            string_data=csv_content,
+            key=latest_file_key,
+            bucket_name=bucket_name,
+            replace=True
+        )
+
+    # Push the latest file name to XCom
+    if 'latest_file_key' in locals():
+        kwargs['ti'].xcom_push(key='file_name', value=latest_file_key)
+    else:
+        print("No files found in the specified path.")
 
 def generate_s3_path(client_name):
     # Your function as provided
@@ -58,6 +74,23 @@ def generate_s3_path(client_name):
     base_s3_path = 's3://'
     return base_s3_path + final_name + "/"
 
+def call_stored_proc(proc_name, **kwargs):
+     # Pull the file name from XCom
+    file_name = kwargs['ti'].xcom_pull(task_ids='process_csv', key='file_name')
+    logger.info(f"File name: {file_name}")
+    dwh_hook =  SnowflakeHook(snowflake_conn_id=SNOWFLAKE_CONN_ID)
+    with dwh_hook.get_conn() as conn:
+        # dwh_hook.set_autocommit(conn,autocommit=False)
+        cur = conn.cursor()
+
+        cur.execute(f"CALL {DATABASE}.STG.{proc_name}('{file_name}');")
+        result = cur.fetchone()
+        if result[0] == 'Setup, Load, and Audit Complete':
+            logger.info('PROCEDURE EXECUTED SUCCESSFULLY')
+        else:
+            raise AirflowFailException("Check the DAG logs for more information. ERROR FROM SNOWFLAKE: ", result)
+        logger.info(f"QUERY EXECUTION RESULT: {str(result)}")
+
 default_args = {
     'owner': 'airflow',
     'depends_on_past': False,
@@ -70,14 +103,29 @@ dag = DAG(
     DAG_ID,
     default_args=default_args,
     description='Process a CSV in S3 and replace it',
-    schedule_interval='@daily',
-    params = default_params
+    schedule_interval='@daily'
 )
 
+begin_job = EmptyOperator(task_id='Begin')
+
 process_csv_task = PythonOperator(
     task_id='process_csv',
     python_callable=process_csv_in_s3,
+    provide_context=True,
     dag=dag
 )
 
-process_csv_task
+load_client_config = PythonOperator(
+    task_id="load_raw_training_data",
+    python_callable=call_stored_proc,
+    dag=dag,
+    provide_context=True,
+    op_kwargs={'proc_name':'LOAD_CLIENT_CONFIG'}        
+)
+
+
+
+
+end_job = EmptyOperator(task_id='End')
+
+begin_job >> process_csv_task >> load_client_config >> end_job