Added UI 3 dag (tested and deployed)

2024-02-29 12:45:05 -06:00
parent 5dbd7d0072
commit 33806046d1
1 changed files with 89 additions and 0 deletions
@@ -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