Added QC report dag. Tested and Deployed
This commit is contained in:
@@ -0,0 +1,81 @@
|
||||
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_dev_snowflake"
|
||||
TAGS=["QC","dev","etl","QC-Summery"]
|
||||
DAG_ID="qc_report_dag"
|
||||
|
||||
bucket = "doczy-dev-infra-mwaa-resources"
|
||||
object_key = "outputs/"
|
||||
|
||||
DATABASE="DOCZY_DEV"
|
||||
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}")
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user