From 4da2a706591c3f77a322a6339b06b471f2af81df Mon Sep 17 00:00:00 2001
From: Umang Mistry <umistry@aarete.com>
Date: Thu, 29 Feb 2024 16:05:28 -0600
Subject: [PATCH] Added QC report dag. Tested and Deployed

---
 airflow/Qa_Dag.py | 81 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)
 create mode 100644 airflow/Qa_Dag.py

diff --git a/airflow/Qa_Dag.py b/airflow/Qa_Dag.py
new file mode 100644
index 0000000..bceed96
--- /dev/null
+++ b/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
\ No newline at end of file