Notes

import boto3
import gzip
from sqlalchemy import create_engine, text
import urllib3
import logging

urllib3.disable_warnings()

# AWS S3 configuration
S3_BUCKET_NAME = 'aa00002695-quanthub-market-data-s3-dev-us-east-1'
S3_PATH = 'Datastore/Xpressfeed/compustatmarketcap/'
S3_CLIENT = boto3.client('s3', verify=False)

# Database configuration
DB_USER = 'consumg'
DB_PASSWORD = 'consumg#8894'
DB_HOST = 'aa00002695-qhub-prd-clus.cluster-c0wjzrxtrwvo.us-east-1.rds.amazonaws.com'
DB_PORT = '5440'
DB_NAME = 'qhub'
SCHEMA_NAME = 'consumg'
FINAL_TABLE_NAME = 'compustatmarketcap'

# Database connection string
DB_URL = f'postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}'
ENGINE = create_engine(DB_URL)

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

def list_files(bucket_name, prefix, suffixes=('.csv.gz', '.csv.zip')):
"""List files in S3 bucket with given prefix and suffixes."""
logging.info("Fetching list of files from S3")
try:
response = S3_CLIENT.list_objects_v2(Bucket=bucket_name, Prefix=prefix)
files = [obj['Key'] for obj in response.get('Contents', []) if obj['Key'].endswith(suffixes)]
return files
except Exception as e:
logging.error(f"Error listing files from S3: {str(e)}")
return []

def create_temp_table(engine, schema_name, table_name):
"""Create temporary table in the database."""
try:
with engine.connect() as connection:
logging.info("Database connection successful")
drop_statement = text(f"DROP TABLE IF EXISTS {schema_name}.{table_name}")
connection.execute(drop_statement)
connection.commit()
logging.info(f"Temporary table {table_name} dropped if it existed")
sql_statement = text(f"""
CREATE TABLE {schema_name}.{table_name} (
GVKEY VARCHAR(6),
DATADATE DATE,
CURCDD VARCHAR,
MKVAL DOUBLE PRECISION
)
""")
connection.execute(sql_statement)
connection.commit()
logging.info("Temporary table created successfully")
except Exception as e:
logging.error(f"Error creating temporary table: {str(e)}")
raise

def copy_data_to_temp_table(engine, schema_name, table_name, files):
"""Copy data from S3 files to temporary database table using COPY command."""
try:
logging.info("Starting to copy data to the database table")
for file_key in files:
logging.info(f"Copying data from {file_key}")
try:
obj = S3_CLIENT.get_object(Bucket=S3_BUCKET_NAME, Key=file_key)
with gzip.open(obj['Body'], 'rt') as gz:
copy_sql = f"COPY {schema_name}.{table_name} FROM stdin WITH CSV HEADER DELIMITER ','"
with engine.connect() as conn:
conn.connection.set_isolation_level(0)
with conn.connection.cursor() as cur:
cur.copy_expert(sql=copy_sql, file=gz)
conn.connection.set_isolation_level(1)
cur.execute(f"SELECT COUNT(*) FROM {schema_name}.{table_name}")
count = cur.fetchone()[0]
logging.info(f"Total records copied: {count}")
except Exception as e:
logging.error(f"Failed to copy data from {file_key}: {str(e)}")
raise # Raise exception if copying fails
return True
except Exception as e:
logging.error(f"Error copying data to temporary table: {str(e)}")
return False

def truncate_table(engine, schema_name, table_name):
"""Truncate the final table in the database."""
try:
with engine.connect() as connection:
logging.info(f"Attempting to truncate table {schema_name}.{table_name}")
truncate_statement = text(f"TRUNCATE TABLE {schema_name}.{table_name}")
connection.execute(truncate_statement)
connection.commit()
logging.info(f"Table {table_name} truncated successfully")
except Exception as e:
logging.error(f"Error truncating table {table_name}: {str(e)}")
raise

def rename_and_replace_table(engine, schema_name, temp_table_name, final_table_name):
"""Rename temporary table and replace final table."""
try:
with engine.connect() as connection:
logging.info(f"Attempting to rename table {schema_name}.{temp_table_name} to {final_table_name}")
rename_statement = text(f"ALTER TABLE {schema_name}.{temp_table_name} RENAME TO {final_table_name}")
connection.execute(rename_statement)
connection.commit()
logging.info(f"Temporary table renamed to {final_table_name}")
except Exception as e:
logging.error(f"Error renaming and replacing table: {str(e)}")
raise

def main():
logging.info("Starting the file listing process")
files = list_files(S3_BUCKET_NAME, S3_PATH)
logging.info(f"Files: {files}")
if files:
files = files[:5] # Limit to the first 5 files
temp_table_name = 'temp_table'
try:
create_temp_table(ENGINE, SCHEMA_NAME, temp_table_name)
success = copy_data_to_temp_table(ENGINE, SCHEMA_NAME, temp_table_name, files)
if success:
try:
truncate_table(ENGINE, SCHEMA_NAME, FINAL_TABLE_NAME) # Truncate final table after successful data copy
rename_and_replace_table(ENGINE, SCHEMA_NAME, temp_table_name, FINAL_TABLE_NAME)
except Exception as e:
logging.error(f"Error during final table update: {str(e)}")
raise
else:
logging.error("Copying data to temporary table failed. Skipping truncating and renaming.")
except Exception as e:
logging.error(f"Process failed: {str(e)}. Final table remains unchanged.")
else:
logging.info("No CSV files found")

if __name__ == "__main__":
main()