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 diff --git a/airflow/training_results_dag.py b/airflow/training_results_dag.py new file mode 100644 index 0000000..40cad67 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/snowflake/scripts/training_interface/R__002_LOAD_TRAINING_RESULTS_SP.SQL b/snowflake/scripts/training_interface/R__002_LOAD_TRAINING_RESULTS_SP.SQL index ea3b3cf..e3f20af 100644 --- a/snowflake/scripts/training_interface/R__002_LOAD_TRAINING_RESULTS_SP.SQL +++ b/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'); diff --git a/snowflake/scripts/training_interface/R__003_LOAD_ATTEMPTS_SP.sql b/snowflake/scripts/training_interface/R__003_LOAD_ATTEMPTS_SP.sql index 1c7127b..3699dcc 100644 --- a/snowflake/scripts/training_interface/R__003_LOAD_ATTEMPTS_SP.sql +++ b/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');