diff --git a/airflow/dags/uat_Qa_Dag.py b/airflow/dags/uat_Qa_Dag.py new file mode 100644 index 0000000..bfda702 --- /dev/null +++ b/airflow/dags/uat_Qa_Dag.py @@ -0,0 +1,94 @@ +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_uat_snowflake" +TAGS=["QC","uat","etl","QC-Summery"] +DAG_ID="qc_report_dag" + +bucket = "doczy-dev-infra-mwaa-resources" +object_key = "outputs/" + +DATABASE="DOCZY_UAT" +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}") + + oldData = df + newData = df + + with io.BytesIO() as output: + with pd.ExcelWriter(output, engine='xlsxwriter') as writer: + oldData.to_excel(writer, sheet_name="Old") + newData.to_excel(writer, sheet_name="new") + dat = oldData.compare(newData, align_axis=0, keep_shape=True) + dat.to_excel(writer, sheet_name="qc") + response = s3.put_object( + Bucket=bucket, Key=object_key+'QC_Report.xlsx' , Body=output.getvalue() + ) + + 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/dags/uat_config_interface_dag .py b/airflow/dags/uat_config_interface_dag .py new file mode 100644 index 0000000..bcca34f --- /dev/null +++ b/airflow/dags/uat_config_interface_dag .py @@ -0,0 +1,100 @@ + + +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_uat_snowflake" +DAG_ID = "load_request_and_contract_submissions" +DATABASE="DOCZY_UAT" +# bucket = "airflow-data-ingestion" + +TAGS=["uat","contract_config_interface","dataload"] + +# Trigger rules +ALL_SUCCESS = 'all_success' +ALL_FAILED = 'all_failed' +ALL_DONE = 'all_done' +ONE_SUCCESS = 'one_success' +ONE_FAILED = 'one_failed' + +# Passing empty params for now, this will be overridden by the payload from the trigger +# These params can also be set from the Airflow UI while manually triggering the DAG +default_params = {"request_submission_file_name": "", "contract_config_file_name":""} +# 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_type, params): + 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() + + # Added new parameter file_type to determine the file name to be passed to the stored procedure + # The bucket name is set by default to "doczy-dev-infra-raw-data-ingestion" and the files should be ALWAYS save under training_interface/ path for now + if file_type == 'request_submission': + file_name = params['request_submission_file_name'] + elif file_type == 'contract_config': + file_name = params['contract_config_file_name'] + + 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, + params = default_params +) + +begin_job = EmptyOperator(task_id='Begin') + + +load_request_submission = PythonOperator( + task_id="load_request_submissions", + python_callable=call_stored_proc, + dag=dag, + op_kwargs={'proc_name':'LOAD_REQUEST_SUBMISSION', 'file_type':'request_submission'} +) + +load_contract_config = PythonOperator( + task_id="load_contract_config", + python_callable=call_stored_proc, + dag=dag, + op_kwargs={'proc_name':'LOAD_CONTRACT_CONFIG', 'file_type':'contract_config'} +) + + + + +end_job = EmptyOperator(task_id='End') + + +begin_job >> load_request_submission >> load_contract_config >> end_job \ No newline at end of file diff --git a/airflow/dags/uat_raw_training_data_dag.py b/airflow/dags/uat_raw_training_data_dag.py new file mode 100644 index 0000000..c78fda5 --- /dev/null +++ b/airflow/dags/uat_raw_training_data_dag.py @@ -0,0 +1,108 @@ + + +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_uat_snowflake" +DAG_ID = "load_raw_training_data" +DATABASE="DOCZY_UAT" +# bucket = "airflow-data-ingestion" + +TAGS=["uat","training_data","dataload"] + +# Trigger rules +ALL_SUCCESS = 'all_success' +ALL_FAILED = 'all_failed' +ALL_DONE = 'all_done' +ONE_SUCCESS = 'one_success' +ONE_FAILED = 'one_failed' + +# Passing empty params for now, this will be overridden by the payload from the trigger +# These params can also be set from the Airflow UI while manually triggering the DAG +default_params = {"column_config_file_name": "", "raw_training_data_file_name":"", "business_config_file_name":""} +# 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_type, params): + 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() + + # Added new parameter file_type to determine the file name to be passed to the stored procedure + # The bucket name is set by default to "doczy-dev-infra-raw-data-ingestion" and the files should be ALWAYS save under training_interface/ path for now + if file_type == 'column_config': + file_name = params['column_config_file_name'] + elif file_type == 'raw_training_data': + file_name = params['raw_training_data_file_name'] + elif file_type == 'business_config': + file_name = params['business_config_file_name'] + + 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, + params = default_params +) + +begin_job = EmptyOperator(task_id='Begin') + + +load_training_results = PythonOperator( + task_id="load_column_config", + python_callable=call_stored_proc, + dag=dag, + op_kwargs={'proc_name':'LOAD_COLUMN_CONFIG', 'file_type':'column_config'} +) + +load_attempt_logs_sp = PythonOperator( + task_id="load_raw_training_data", + python_callable=call_stored_proc, + dag=dag, + op_kwargs={'proc_name':'LOAD_TRAINING_DATA_RAW', 'file_type':'raw_training_data'} +) + +load_business_config = PythonOperator( + task_id="load_business_config", + python_callable=call_stored_proc, + dag=dag, + op_kwargs={'proc_name':'LOAD_BUSINESS_CONFIG', 'file_type':'business_config'} +) + + + +end_job = EmptyOperator(task_id='End') + + +begin_job >> load_training_results >> load_attempt_logs_sp >> load_business_config >> end_job \ No newline at end of file