Umang Mistry
84 lines
2.2 KiB
Python
84 lines
2.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
|
|
|
|
|
|
SNOWFLAKE_CONN_ID = "doczy_dev_snowflake"
|
|
DAG_ID = "load_client_config"
|
|
DATABASE="DOCZY_DEV"
|
|
# bucket = "airflow-data-ingestion"
|
|
|
|
TAGS=["dev","config_interface","dataload"]
|
|
|
|
default_params = {"client_list_file_name": ""}
|
|
|
|
|
|
|
|
def process_csv_in_s3(params):
|
|
# Instantiate S3Hook
|
|
file_name = params['client_list_file_name']
|
|
s3_hook = S3Hook(aws_conn_id="aws_default")
|
|
|
|
bucket_name = "doczy-dev-infra-raw-data-ingestion"
|
|
key = 'client_names_openair/' + file_name
|
|
# Get the file from S3
|
|
file_content = s3_hook.read_key(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=key,
|
|
bucket_name=bucket_name,
|
|
replace=True
|
|
)
|
|
|
|
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 + "/"
|
|
|
|
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,
|
|
description='Process a CSV in S3 and replace it',
|
|
schedule_interval='@daily',
|
|
params = default_params
|
|
)
|
|
|
|
process_csv_task = PythonOperator(
|
|
task_id='process_csv',
|
|
python_callable=process_csv_in_s3,
|
|
dag=dag
|
|
)
|
|
|
|
process_csv_task
|