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