doczyai-pipelines/archive/airflow/dags/raw_training_data_dag.py

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
)