Created UAT data DAGs in airflow
This commit is contained in:
Umang Mistry
@@ -0,0 +1,94 @@
|
||||
import logging
|
||||
from airflow.exceptions import AirflowFailException
|
||||
from airflow.operators.empty import EmptyOperator
|
||||
from airflow.operators.python import PythonOperator
|
||||
from datetime import datetime
|
||||
from airflow import DAG
|
||||
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
|
||||
# from openpyxl.workbook import Workbook
|
||||
import pandas as pd
|
||||
import boto3
|
||||
import io
|
||||
|
||||
|
||||
SNOWFLAKE_CONN_ID="doczy_uat_snowflake"
|
||||
TAGS=["QC","uat","etl","QC-Summery"]
|
||||
DAG_ID="qc_report_dag"
|
||||
|
||||
bucket = "doczy-dev-infra-mwaa-resources"
|
||||
object_key = "outputs/"
|
||||
|
||||
DATABASE="DOCZY_UAT"
|
||||
SCHEMA = "STG"
|
||||
TABLE_NAME = "TRAINING_DATA_RAW"
|
||||
|
||||
# Trigger rules
|
||||
ALL_SUCCESS = 'all_success'
|
||||
ALL_FAILED = 'all_failed'
|
||||
ALL_DONE = 'all_done'
|
||||
ONE_SUCCESS = 'one_success'
|
||||
ONE_FAILED = 'one_failed'
|
||||
|
||||
|
||||
|
||||
|
||||
args = {"owner": "Airflow", "start_date": datetime(2022, 1, 1), "retries":0 }
|
||||
dag = DAG(
|
||||
dag_id=DAG_ID, default_args=args, schedule=None,
|
||||
tags=TAGS
|
||||
)
|
||||
|
||||
def getData():
|
||||
# Setup connection to Snowflake
|
||||
dwh_hook = SnowflakeHook(snowflake_conn_id=SNOWFLAKE_CONN_ID)
|
||||
conn = dwh_hook.get_conn() # Get the raw connection
|
||||
|
||||
# Your query and the database details
|
||||
qc_query = f"SELECT * FROM {DATABASE}.{SCHEMA}.{TABLE_NAME}"
|
||||
|
||||
# Fetch data into a Pandas DataFrame
|
||||
df = pd.read_sql(qc_query, conn)
|
||||
|
||||
# Initialize S3 client
|
||||
s3 = boto3.client("s3")
|
||||
|
||||
# Create a buffer to hold the data
|
||||
with io.StringIO() as csv_buffer:
|
||||
df.to_csv(csv_buffer, index=False)
|
||||
|
||||
# Save the data to S3
|
||||
response = s3.put_object(
|
||||
Bucket=bucket, Key=object_key+'snowflake_table_results.csv' , Body=csv_buffer.getvalue()
|
||||
)
|
||||
|
||||
status = response.get("ResponseMetadata", {}).get("HTTPStatusCode")
|
||||
|
||||
if status == 200:
|
||||
print(f"Successful S3 put_object response. Status - {status}")
|
||||
else:
|
||||
raise AirflowFailException(f"Unsuccessful S3 put_object response. Status - {status}")
|
||||
|
||||
oldData = df
|
||||
newData = df
|
||||
|
||||
with io.BytesIO() as output:
|
||||
with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
|
||||
oldData.to_excel(writer, sheet_name="Old")
|
||||
newData.to_excel(writer, sheet_name="new")
|
||||
dat = oldData.compare(newData, align_axis=0, keep_shape=True)
|
||||
dat.to_excel(writer, sheet_name="qc")
|
||||
response = s3.put_object(
|
||||
Bucket=bucket, Key=object_key+'QC_Report.xlsx' , Body=output.getvalue()
|
||||
)
|
||||
|
||||
print("Dataframe is written to S3 successfully.")
|
||||
|
||||
with dag:
|
||||
begin_job = EmptyOperator(task_id='Begin')
|
||||
|
||||
get_data_from_snowflake = PythonOperator(task_id="get_data_from_snowflake", python_callable=getData)
|
||||
|
||||
end_job = EmptyOperator(task_id='End')
|
||||
|
||||
|
||||
begin_job >> get_data_from_snowflake >> end_job
|
||||
@@ -0,0 +1,100 @@
|
||||
|
||||
|
||||
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_uat_snowflake"
|
||||
DAG_ID = "load_request_and_contract_submissions"
|
||||
DATABASE="DOCZY_UAT"
|
||||
# bucket = "airflow-data-ingestion"
|
||||
|
||||
TAGS=["uat","contract_config_interface","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 = {"request_submission_file_name": "", "contract_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 == 'request_submission':
|
||||
file_name = params['request_submission_file_name']
|
||||
elif file_type == 'contract_config':
|
||||
file_name = params['contract_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_request_submission = PythonOperator(
|
||||
task_id="load_request_submissions",
|
||||
python_callable=call_stored_proc,
|
||||
dag=dag,
|
||||
op_kwargs={'proc_name':'LOAD_REQUEST_SUBMISSION', 'file_type':'request_submission'}
|
||||
)
|
||||
|
||||
load_contract_config = PythonOperator(
|
||||
task_id="load_contract_config",
|
||||
python_callable=call_stored_proc,
|
||||
dag=dag,
|
||||
op_kwargs={'proc_name':'LOAD_CONTRACT_CONFIG', 'file_type':'contract_config'}
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
end_job = EmptyOperator(task_id='End')
|
||||
|
||||
|
||||
begin_job >> load_request_submission >> load_contract_config >> end_job
|
||||
@@ -0,0 +1,108 @@
|
||||
|
||||
|
||||
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_uat_snowflake"
|
||||
DAG_ID = "load_raw_training_data"
|
||||
DATABASE="DOCZY_UAT"
|
||||
# bucket = "airflow-data-ingestion"
|
||||
|
||||
TAGS=["uat","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
|
||||
Reference in New Issue
Block a user