Appendix A: Python Scripts

This section provides the python scripts used in our use cases.

A.1. PostgreSQL to Amazon S3

import psycopg2
import pandas as pd
import boto3
from io import BytesIO
from datetime import datetime, timedelta
import logging

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler()]
)

def extract_data_from_postgres(conn, query):
    try:
        cursor = conn.cursor()
        cursor.execute(query)
        rows = cursor.fetchall()
        columns = [desc[0] for desc in cursor.description]
        df = pd.DataFrame(rows, columns=columns)
        logging.info(f"Data extracted: {len(df)} records.")
        return df
    except Exception as e:
        logging.error(f"Error during data extraction: {e}")
        raise
        return None

def upload_to_s3(df, aws_access_key_id, aws_secret_access_key, s3_bucket_name, s3_prefix, file_label):
    try:
        buffer = BytesIO()
        df.to_parquet(buffer, index=False, engine='pyarrow')
        buffer.seek(0)

        today_str = (datetime.today() + timedelta(days=1)).strftime('%Y-%m-%d')
        timestamp_str = datetime.today().strftime('%Y%m%d_%H%M%S')

        folder_path = f"{s3_prefix}/{file_label}/insertion_date={today_str}/"
        file_name = f"{file_label}-{timestamp_str}.parquet"
        s3_path = folder_path + file_name

        s3 = boto3.client(
            's3',
            aws_access_key_id=aws_access_key_id,
            aws_secret_access_key=aws_secret_access_key
        )
        s3.upload_fileobj(buffer, s3_bucket_name, s3_path)

        logging.info(f"Uploaded to s3://{s3_bucket_name}/{s3_path}")
    except Exception as e:
        logging.error(f"Error uploading to S3 for {file_label}: {e}")
        raise

def main(aws_access_key_id, aws_secret_access_key, s3_bucket_name, s3_prefix,
         host, port, database, username, password, start_date=None, end_date=None,day=None):
    if day is not None and start_date is None and end_date is None:
        # Calculate 'current date - 1 day'
        current_date = datetime.now().date() # Get only the date part
        calculated_date = current_date - timedelta(days=int(day)) # Use the 'day' parameter here

        start_date = calculated_date
        end_date = calculated_date
    elif start_date is not None and end_date is None:
        # If only start_date is given, assume end_date is the same
        end_date = start_date
    try:
        conn = psycopg2.connect(
            dbname=database,
            user=username,
            password=password,
            host=host,
            port=port
        )
        logging.info("Connected to PostgreSQL.")

        queries = [
            {
                "label": "sale",
                "sql": fr"""  select column1,2,3..... from sale
                            WHERE s.transaction_date::date BETWEEN DATE '{start_date}' AND DATE '{end_date}' """
            },
            {
                "label": "sale_item",
                "sql":fr"""   select ........... """
            },
            {
                "label": "transaction",
                "sql":f"""  ....."""}

            ]

        for query_info in queries:
            label = query_info["label"]
            sql = query_info["sql"]
            logging.info(f"Running query for: {label}")
            df = extract_data_from_postgres(conn, sql)
            if df is None or df.empty:
                logging.warning(f"No data found for: {label}")
                continue
            upload_to_s3(df, aws_access_key_id, aws_secret_access_key, s3_bucket_name, s3_prefix, label)

    except Exception as e:
        logging.error(f"Error in main(): {e}")
        raise
    finally:
        if 'conn' in locals():
            conn.close()
            logging.info("PostgreSQL connection closed.")

A.2. Amazon S3 to PostgreSQL

import boto3
import pandas as pd
import pyarrow.parquet as pq
from io import BytesIO
import logging
import re
from sqlalchemy import create_engine, inspect
import psycopg2
from datetime import datetime

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def create_engine_connection(db_config):
    return create_engine(
        f"postgresql+psycopg2://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}"
    )

def push_dataframe_to_db(df, engine, schema, table_name):
    inspector = inspect(engine)
    table_columns = [col['name'] for col in inspector.get_columns(table_name, schema=schema)]
    filtered_df = df[[col for col in df.columns if col in table_columns]]
    filtered_df.to_sql(table_name, engine, schema=schema, if_exists='append', index=False, method='multi')
    return len(filtered_df)

def create_combined_dataframe_from_latest_partition(bucket, parent_prefix, aws_key, aws_secret, filter_tables=None):
    s3 = boto3.client('s3', aws_access_key_id=aws_key, aws_secret_access_key=aws_secret)
    response = s3.list_objects_v2(Bucket=bucket, Prefix=parent_prefix, Delimiter='/')
    subfolders = [prefix['Prefix'] for prefix in response.get('CommonPrefixes', [])]
    table_dataframes = {}

    for subfolder in subfolders:
        table_name = subfolder.replace(parent_prefix, '').strip('/')
        if filter_tables and table_name not in filter_tables:
            continue

        date_folders_resp = s3.list_objects_v2(Bucket=bucket, Prefix=subfolder, Delimiter='/')
        partition_folders = []
        for prefix in date_folders_resp.get('CommonPrefixes', []):
            folder_key = prefix['Prefix']
            match = re.search(r'insertion_date=(\d{4}-\d{2}-\d{2})/', folder_key)
            if match:
                partition_folders.append((match.group(1), folder_key))

        if not partition_folders:
            continue

        latest_partition_folder = max(partition_folders, key=lambda x: x[0])[1]
        file_resp = s3.list_objects_v2(Bucket=bucket, Prefix=latest_partition_folder)
        dfs = []

        for obj in file_resp.get('Contents', []):
            key = obj['Key']
            if key.endswith('.parquet'):
                parquet_obj = s3.get_object(Bucket=bucket, Key=key)
                df = pq.read_table(source=BytesIO(parquet_obj['Body'].read())).to_pandas()
                dfs.append(df)

        if dfs:
            combined_df = pd.concat(dfs, ignore_index=True)
            table_dataframes[table_name] = combined_df

    return table_dataframes

def main(
    aws_access_key_id, aws_secret_access_key,
    s3_bucket_name, s3_prefix,
    host, port, database, username, password,
    filter_tables=None,
    schema="public"
):
    start_time = datetime.now()
    process_id = None
    tables_processed = 0
    tables_successful = 0
    tables_failed = 0

    try:
        logging.info("Starting ETL job...")

        db_config = {
            'host': host,
            'port': port,
            'database': database,
            'user': username,
            'password': password
        }

        engine = create_engine_connection(db_config)
        audit_conn = psycopg2.connect(
            host=host, port=port, dbname=database, user=username, password=password
        )

        with audit_conn.cursor() as cur:
            cur.execute("""
                INSERT INTO silver_etl_master.silver_master_etl_audit_tbl (
                    schedule_name, schedule_timestamp, status, start_time,
                    tables_processed, tables_failed, tables_successful
                ) VALUES (%s, %s, %s, %s, %s, %s, %s)
                RETURNING process_id
            """, (
                's3_partitioned_etl',
                datetime.now(),
                'In Progress',
                start_time,
                0, 0, 0
            ))
            process_id = cur.fetchone()[0]
            audit_conn.commit()

        table_dfs = create_combined_dataframe_from_latest_partition(
            s3_bucket_name, s3_prefix,
            aws_access_key_id, aws_secret_access_key,
            filter_tables=filter_tables
        )

        if not table_dfs:
            raise Exception("No dataframes found in S3.")

        for table_name, df in table_dfs.items():
            tables_processed += 1
            try:
                inserted = push_dataframe_to_db(df, engine, schema, table_name)

                with audit_conn.cursor() as cur:
                    cur.execute("""
                        INSERT INTO silver_etl_master.silver_main_etl_audit_tbl (
                            process_id, data_source, source_table_name, destination_table_name,
                            operation, etl_start_time, etl_end_time, input_row, output_row, operation_result
                        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    """, (
                        process_id,
                        'S3',
                        [f"{s3_prefix}/{table_name}"],
                        f"{schema}.{table_name}",
                        'insert',
                        start_time,
                        datetime.now(),
                        len(df),
                        inserted,
                        'Success'
                    ))
                audit_conn.commit()
                tables_successful += 1

            except Exception as e:
                logging.error(f"Failed to process table {table_name}: {e}")
                with audit_conn.cursor() as cur:
                    cur.execute("""
                        INSERT INTO silver_etl_master.silver_main_etl_audit_tbl (
                            process_id, data_source, source_table_name, destination_table_name,
                            operation, etl_start_time, etl_end_time, input_row, output_row,
                            operation_result, result
                        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    """, (
                        process_id,
                        'S3',
                        [f"{s3_prefix}/{table_name}"],
                        f"{schema}.{table_name}",
                        'insert',
                        start_time,
                        datetime.now(),
                        len(df) if df is not None else 0,
                        0,
                        'Failed',
                        str(e)
                    ))
                audit_conn.commit()
                tables_failed += 1

        # Final update of master ETL table
        with audit_conn.cursor() as cur:
            cur.execute("""
                UPDATE silver_etl_master.silver_master_etl_audit_tbl
                SET status = %s, end_time = %s,
                    tables_processed = %s, tables_failed = %s, tables_successful = %s
                WHERE process_id = %s
            """, (
                'Completed',
                datetime.now(),
                tables_processed,
                tables_failed,
                tables_successful,
                process_id
            ))
            audit_conn.commit()

    except Exception as e:
        logging.error(f"ETL job failed: {e}")
        if process_id:
            with psycopg2.connect(
                host=host, port=port, dbname=database, user=username, password=password
            ) as audit_conn:
                with audit_conn.cursor() as cur:
                    cur.execute("""
                        UPDATE silver_etl_master.silver_master_etl_audit_tbl
                        SET status = %s, end_time = %s, error_message = %s
                        WHERE process_id = %s
                    """, ('Failed', datetime.now(), str(e), process_id))
                    audit_conn.commit()
    finally:
        logging.info("ETL job completed.")