Michael McGuinness
119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
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_raw_training_data"
|
|
DATABASE = "DOCZY_DEV"
|
|
# bucket = "airflow-data-ingestion"
|
|
|
|
TAGS = ["dev", "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
|
|
)
|