Notes

import boto3
import gzip
import pandas as pd
from sqlalchemy import create_engine, text

# 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)
batch_size = 100000 # Adjust based on memory availability

# Database configuration
db_user = 'your_db_user'
db_password = 'your_db_password'
db_host = 'your_db_host'
db_port = 'your_db_port'
db_name = 'your_db_name'
schema_name = 'your_schema_name'
temp_table_name = 'temp_table'
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)

def list_files(bucket_name, prefix, suffix='.csv.gz'):
print("Fetching list")
response = S3_client.list_objects_v2(Bucket=bucket_name, Prefix=prefix)
if 'Contents' in response:
files = [obj['Key'] for obj in response['Contents'] if obj['Key'].endswith(suffix)]
return files
else:
return []

def create_temp_table(engine, schema_name, table_name):
with engine.connect() as connection:
print("Database connection successful") # Print check
drop_statement = text(f"DROP TABLE IF EXISTS {schema_name}.{table_name}")
connection.execute(drop_statement)
print(f"Temporary table {table_name} dropped if it existed") # Print check
sql_statement = text(f"CREATE TEMP TABLE {schema_name}.{table_name} (GVKEY VARCHAR(6), DATADATE DATE, CURCDD VARCHAR, MKVAL DOUBLE PRECISION)")
connection.execute(sql_statement)
print("Temporary table created successfully") # Print check

def load_data_to_temp_table(engine, schema_name, table_name, file_key, batch_size):
print(f"Starting to load data from {file_key} to the database table") # Print check
obj = S3_client.get_object(Bucket=S3_bucket_name, Key=file_key)
with gzip.open(obj['Body'], 'rt') as gz:
reader = pd.read_csv(gz, chunksize=batch_size)
for chunk in reader:
chunk.to_sql(table_name, engine, schema=schema_name, if_exists='append', index=False)

def rename_and_replace_table(engine, schema_name, temp_table_name, final_table_name):
with engine.connect() as connection:
# Drop the final table if it exists
drop_statement = text(f"DROP TABLE IF EXISTS {schema_name}.{final_table_name}")
connection.execute(drop_statement)
print(f"Existing table {final_table_name} dropped") # Print check

# Rename the temporary table to the final table name
rename_statement = text(f"ALTER TABLE {schema_name}.{temp_table_name} RENAME TO {final_table_name}")
connection.execute(rename_statement)
print(f"Temporary table renamed to {final_table_name}") # Print check

def main():
print("Starting the file listing process")
files = list_files(S3_bucket_name, s3_path)
print("Files:")
if files:
print(f"CSV.gz files in {S3_bucket_name}/{s3_path}:")
for file_key in files:
print(file_key)

create_temp_table(engine, schema_name, temp_table_name)

# Load data from all files into the temporary table
for file_key in files:
load_data_to_temp_table(engine, schema_name, temp_table_name, file_key, batch_size)

# Rename and replace the final table
rename_and_replace_table(engine, schema_name, temp_table_name, final_table_name)

else:
print("No CSV files found")

if __name__ == "__main__":
main()