Umang Mistry
134 lines
4.2 KiB
Python
134 lines
4.2 KiB
Python
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
|