From 1fa1b8883fcc0c4f5366864c4c01f2f1d1f454f2 Mon Sep 17 00:00:00 2001 From: Umang Mistry Date: Thu, 7 Mar 2024 17:17:49 -0600 Subject: [PATCH] Updated UI pipeline with parameter that will be passed by UI --- airflow/dags/cicd_test_dag.py | 2 +- airflow/dags/training_results_dag.py | 23 +++++++++++++++++------ 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/airflow/dags/cicd_test_dag.py b/airflow/dags/cicd_test_dag.py index c9f63fa..20e7351 100644 --- a/airflow/dags/cicd_test_dag.py +++ b/airflow/dags/cicd_test_dag.py @@ -137,7 +137,7 @@ begin_job = EmptyOperator(task_id='Begin') python_task = PythonOperator(task_id='get_info_using_python_op', # Task ID has to be unique only inside a dags python_callable=python_op_eg, dag=dag, - op_kwargs={'table_name':'DIM_AUDIT'} # Variables can be passed to the python function using op_kwargs + op_kwargs={'table_name':'DIM_AUDIT',} # Variables can be passed to the python function using op_kwargs ) # Snowflake operator does not offer the option to display query results diff --git a/airflow/dags/training_results_dag.py b/airflow/dags/training_results_dag.py index 40cad67..7c01caa 100644 --- a/airflow/dags/training_results_dag.py +++ b/airflow/dags/training_results_dag.py @@ -35,17 +35,27 @@ 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 = {"training_results_file_name": "", "attempt_logs_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" +# training_results_file_name = "training_results_sample.csv" +# attempt_logs_file_name = "attempt_logs_sample.csv" -def call_stored_proc(proc_name,file_name): +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 == 'training_results': + file_name = params['training_results_file_name'] + elif file_type == 'attempt_logs': + file_name = params['attempt_logs_file_name'] + cur.execute(f"CALL {DATABASE}.STG.{proc_name}('{file_name}');") result = cur.fetchone() if result[0] == 'Setup, Load, and Audit Complete': @@ -60,7 +70,8 @@ dag = DAG( default_args={"snowflake_conn_id": SNOWFLAKE_CONN_ID, "retries":0}, tags=TAGS, catchup=False, - schedule=None + schedule=None, + params = default_params ) begin_job = EmptyOperator(task_id='Begin') @@ -70,14 +81,14 @@ 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} + op_kwargs={'proc_name':'LOAD_TRAINING_RESULTS', 'file_type':'training_results'} ) 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} + op_kwargs={'proc_name':'LOAD_ATTEMPT_LOGS', 'file_type':'attempt_logs'} )