Michael McGuinness
140 lines
4.2 KiB
Python
140 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
|