Merged in dev_umistry (pull request #13)

Dev umistry

airflow/Qa_Dag.py

@@ -0,0 +1,81 @@
+import logging
+from airflow.exceptions import AirflowFailException
+from airflow.operators.empty import EmptyOperator
+from airflow.operators.python import PythonOperator
+from datetime import datetime
+from airflow import DAG
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+# from openpyxl.workbook import Workbook
+import pandas as pd
+import boto3
+import io
+
+
+SNOWFLAKE_CONN_ID="doczy_dev_snowflake"
+TAGS=["QC","dev","etl","QC-Summery"]
+DAG_ID="qc_report_dag"
+
+bucket = "doczy-dev-infra-mwaa-resources"
+object_key = "outputs/"
+
+DATABASE="DOCZY_DEV"        
+SCHEMA = "STG"
+TABLE_NAME = "TRAINING_DATA_RAW"
+
+# Trigger rules
+ALL_SUCCESS = 'all_success'
+ALL_FAILED = 'all_failed'
+ALL_DONE = 'all_done'
+ONE_SUCCESS = 'one_success'
+ONE_FAILED = 'one_failed'
+
+
+
+
+args = {"owner": "Airflow", "start_date": datetime(2022, 1, 1), "retries":0 }
+dag = DAG(
+    dag_id=DAG_ID, default_args=args, schedule=None,
+    tags=TAGS
+)
+
+def getData():
+    # Setup connection to Snowflake
+    dwh_hook = SnowflakeHook(snowflake_conn_id=SNOWFLAKE_CONN_ID)
+    conn = dwh_hook.get_conn()  # Get the raw connection
+    
+    # Your query and the database details
+    qc_query = f"SELECT * FROM {DATABASE}.{SCHEMA}.{TABLE_NAME}"
+    
+    # Fetch data into a Pandas DataFrame
+    df = pd.read_sql(qc_query, conn)
+
+    # Initialize S3 client
+    s3 = boto3.client("s3")
+
+    # Create a buffer to hold the data
+    with io.StringIO() as csv_buffer:
+        df.to_csv(csv_buffer, index=False)
+        
+        # Save the data to S3
+        response = s3.put_object(
+            Bucket=bucket, Key=object_key+'snowflake_table_results.csv' , Body=csv_buffer.getvalue()
+        )
+
+        status = response.get("ResponseMetadata", {}).get("HTTPStatusCode")
+
+        if status == 200:
+            print(f"Successful S3 put_object response. Status - {status}")
+        else:
+            raise AirflowFailException(f"Unsuccessful S3 put_object response. Status - {status}")
+
+    print("Dataframe is written to S3 successfully.")
+
+with dag:
+    begin_job = EmptyOperator(task_id='Begin')
+
+    get_data_from_snowflake = PythonOperator(task_id="get_data_from_snowflake", python_callable=getData)
+
+    end_job = EmptyOperator(task_id='End')
+
+
+begin_job >> get_data_from_snowflake >> end_job

airflow/training_results_dag.py

@@ -0,0 +1,89 @@
+
+
+from __future__ import annotations
+
+import os
+from datetime import datetime
+import logging
+
+from airflow import DAG
+from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+from airflow.operators.python import PythonOperator
+from airflow.utils.trigger_rule import TriggerRule
+from airflow.operators.empty import EmptyOperator
+from airflow.models import Variable
+from airflow.exceptions import AirflowFailException
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+
+SNOWFLAKE_CONN_ID = "doczy_dev_snowflake"
+DAG_ID = "load_training_results"
+DATABASE="DOCZY_DEV"
+# bucket = "airflow-data-ingestion"
+
+TAGS=["dev","training_interface","dataload"]
+
+# Trigger rules
+ALL_SUCCESS = 'all_success'
+ALL_FAILED = 'all_failed'
+ALL_DONE = 'all_done'
+ONE_SUCCESS = 'one_success'
+ONE_FAILED = 'one_failed'
+
+# This will be replaced with the payload from the event after API connection is setup
+training_results_file_name = "training_results_sample.csv"
+attempt_logs_file_name = "attempt_logs_sample.csv"
+
+
+def call_stored_proc(proc_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)}")
+
+dag = DAG(
+    DAG_ID,
+    start_date=datetime(2024, 1, 1),
+    default_args={"snowflake_conn_id": SNOWFLAKE_CONN_ID, "retries":0},
+    tags=TAGS,
+    catchup=False,
+    schedule=None
+)
+    
+begin_job = EmptyOperator(task_id='Begin')
+
+
+load_training_results = PythonOperator(
+    task_id="load_training_results",
+    python_callable=call_stored_proc,
+    dag=dag,
+    op_kwargs={'proc_name':'LOAD_TRAINING_RESULTS', 'file_name':training_results_file_name}        
+)
+
+load_attempt_logs_sp = PythonOperator(
+    task_id="load_attempt_logs",
+    python_callable=call_stored_proc,
+    dag=dag,
+    op_kwargs={'proc_name':'LOAD_ATTEMPT_LOGS', 'file_name':attempt_logs_file_name}        
+)
+
+
+
+
+end_job = EmptyOperator(task_id='End')
+
+
+begin_job >> load_training_results >> load_attempt_logs_sp >> end_job

snowflake/scripts/training_interface/R__002_LOAD_TRAINING_RESULTS_SP.SQL

@@ -38,7 +38,7 @@ BEGIN
         FROM @STG.TRAINING_RESULTS_STAGE
     )
     FILE_FORMAT = (FORMAT_NAME = 'STG.CSV_HEADER')
-    ON_ERROR = CONTINUE;
+    ON_ERROR = ABORT_STATEMENT;
 
     call stg.log_audit(:procedure_name, 'Section 2', 99, 'END');
 

snowflake/scripts/training_interface/R__003_LOAD_ATTEMPTS_SP.sql

@@ -35,7 +35,7 @@ BEGIN
         FROM @STG.ATTEMPT_LOGS_STAGE
     )
     FILE_FORMAT = (FORMAT_NAME = 'STG.CSV_HEADER')
-    ON_ERROR = CONTINUE;
+    ON_ERROR = ABORT_STATEMENT;
 
     call stg.log_audit(:procedure_name, 'Section 2', 99, 'END');