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

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}"
            )

    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