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

import re
import pandas as pd
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from datetime import datetime, timedelta
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from airflow.operators.empty import EmptyOperator
import logging
from airflow.exceptions import AirflowFailException


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

SNOWFLAKE_CONN_ID = "doczy_dev_snowflake"
DAG_ID = "load_client_config"
DATABASE = "DOCZY_DEV"
# bucket = "airflow-data-ingestion"

TAGS = ["dev", "config_interface", "dataload"]


def process_csv_in_s3(**kwargs):
    s3_hook = S3Hook(aws_conn_id="aws_default")
    bucket_name = "doczy-dev-infra-raw-data-ingestion"
    key_prefix = "client_names_openair/"

    # List objects within the specified bucket and key prefix
    objects = s3_hook.list_keys(bucket_name=bucket_name, prefix=key_prefix)

    # Sort the objects by their last modified date and select the latest
    if objects:
        latest_file_key = sorted(objects)[
            -1
        ]  # Assuming file names include a timestamp or incrementing number

        # Get the latest file from S3
        file_content = s3_hook.read_key(latest_file_key, bucket_name)

        # Convert string to DataFrame
        from io import StringIO

        df = pd.read_csv(StringIO(file_content))

        # Apply the generate_s3_path function
        df["s3_path"] = df["customer_name"].apply(generate_s3_path)

        # Convert DataFrame to CSV string
        csv_buffer = StringIO()
        df.to_csv(csv_buffer, index=False)
        csv_content = csv_buffer.getvalue()

        # Replace the file in S3 with the processed content
        s3_hook.load_string(
            string_data=csv_content,
            key=latest_file_key,
            bucket_name=bucket_name,
            replace=True,
        )

    # Push the latest file name to XCom
    if "latest_file_key" in locals():
        kwargs["ti"].xcom_push(key="file_name", value=latest_file_key.split("/")[-1])
    else:
        print("No files found in the specified path.")


def generate_s3_path(client_name):
    # Your function as provided
    client_name = client_name.rstrip(".")
    pattern = r"[^0-9a-zA-Z!_.()*\'-]"
    multi_underscore_pattern = r"_{2,}"
    intermediate_name = re.sub(pattern, "_", client_name)
    final_name = re.sub(multi_underscore_pattern, "_", intermediate_name)
    final_name = final_name.lower()
    base_s3_path = "s3://"
    return base_s3_path + final_name + "/"


def call_stored_proc(proc_name, **kwargs):
    # Pull the file name from XCom
    file_name = kwargs["ti"].xcom_pull(task_ids="process_csv", key="file_name")
    logger.info(f"File name: {file_name}")
    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()

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


default_args = {
    "owner": "airflow",
    "depends_on_past": False,
    "start_date": datetime(2024, 3, 28),
    "retries": 1,
    "retry_delay": timedelta(minutes=5),
}

dag = DAG(
    DAG_ID,
    default_args=default_args,
    start_date=datetime(2024, 3, 28),
    catchup=False,
    description="Process a CSV in S3 and replace it",
    schedule_interval="@daily",
)

begin_job = EmptyOperator(task_id="Begin")

process_csv_task = PythonOperator(
    task_id="process_csv",
    python_callable=process_csv_in_s3,
    provide_context=True,
    dag=dag,
)

load_client_config = PythonOperator(
    task_id="load_client_config",
    python_callable=call_stored_proc,
    dag=dag,
    provide_context=True,
    op_kwargs={"proc_name": "LOAD_CLIENT_CONFIG"},
)


end_job = EmptyOperator(task_id="End")

begin_job >> process_csv_task >> load_client_config >> end_job