jacob-mathison/doczyai-pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Updated UI pipeline with parameter that will be passed by UI

Browse Source
This commit is contained in:
Umang Mistry
2024-03-07 17:17:49 -06:00
parent 44c42dcbd0
commit 1fa1b8883f
2 changed files with 18 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
airflow/dags/cicd_test_dag.py
+1 -1
View File
@@ -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
airflow/dags/training_results_dag.py
+17 -6
View File
@@ -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'}
)