doczyai-pipelines/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