add amendments to load lambda

1 files changed, 37 insertions, 29 deletions

diff --git a/src/load_lambda.py b/src/load_lambda.py
index d95c27a..f92bb45 100644
--- a/src/load_lambda.py
+++ b/src/load_lambda.py
@@ -1,11 +1,11 @@
 import boto3
-from botocore.exceptions import ClientError
+from botocore.exceptions import ClientError, InterfaceError
 import pandas as pd
 import pyarrow.parquet as pq
 from io import BytesIO
 import logging
 import json
-from src.extract_lambda import retrieve_secrets, connect_to_database
+from src.extract_lambda import retrieve_secrets
 from sqlalchemy import create_engine
 
 
@@ -18,67 +18,74 @@ logging.basicConfig(
     level=logging.DEBUG,
 )
 
-logging.getLogger("botocore").setLevel(logging.WARNING)
+logging.getLogger("botocore").setLevel(logging.INFO)
+
 
 def lambda_handler(event, context):
-    db = None
     try:
         uploaded_tables = upload_dfs_to_database()
-        if uploaded_tables == []:
+        if not uploaded_tables:
             return {
                 "statusCode": 200,
-                "body": json.dumps("No datframes were uploaded."),
+                "body": json.dumps("No dataframes were uploaded."),
             }
         return {
             "statusCode": 200,
             "body": json.dumps(
                 f"""The following dataframes were uploaded successfully: 
-                {', '.join(upload_dfs_to_database['updated'])}."""
+                {', '.join(uploaded_tables)} ."""
             ),
         }
     except Exception as e:
         logger.error(f"Error: {e}", exc_info=True)
         return {"statusCode": 500, "body": json.dumps("Internal server error.")}
-    finally:
-        if db:
-            db.close()
 
 # connect to database, slightly different way of doing it, to allow manipulation through pandas
 def connect_to_db_and_return_engine():
-    secrets = json.loads(retrieve_secrets("bentley-RDS-credentials"))  #need to amend retrieve secrets function
-    host = secrets["host"]
-    port = secrets["port"]
-    user = secrets["user"]
-    password = secrets["password"]
-    database = secrets["database"]
-    conn_str = f'postgresql+pg8000://{user}:{password}@{host}:{port}/{database}'
-    engine = create_engine(conn_str) #interface between python (pandas) and SQL
-    return engine
-
-
+    try:
+        secrets = json.loads(retrieve_secrets("bentley-RDS-credentials"))  #need to amend retrieve secrets function
+        host = secrets["host"]
+        port = secrets["port"]
+        user = secrets["user"]
+        password = secrets["password"]
+        database = secrets["database"]
+        conn_str = f'postgresql+pg8000://{user}:{password}@{host}:{port}/{database}'
+        engine = create_engine(conn_str) #interface between python (pandas) and SQL
+        return engine
+    except Exception as e:
+        logger.error(f"Interface error: {e}")
+        raise RuntimeError("Failed to create database engine")
+    
 
 # get transform bucket
-def transform_bucket(client=None):
+def get_transform_bucket(client=None):
     if client is None:
         client = boto3.client("s3")
-    response = client.list_buckets()
+    try:
+        response = client.list_buckets()
+    except ClientError as e:
+        logger.error(f"Error listing S3 buckets: {e}")
+        raise RuntimeError("Error listing S3 buckets")
+    
     transform_bucket_filter = [
         bucket["Name"] for bucket in response["Buckets"] if "transform" in bucket["Name"]
     ]
 
     if not transform_bucket_filter:
-        raise ValueError("No transform_bucket found")
+        logger.error("No transform bucket found")
+        raise ValueError("No transform bucket found")
 
     return transform_bucket_filter[0]
 
 # list and then retrieve parquet files from S3 bucket
-# convert parquet files into dataframes and return a list of dataframes  
+# convert parquet files into dataframes
+# return a dictionary of dataframes with name as key, and dataframe object as value 
 def convert_parquet_files_to_dfs(bucket_name=None, client=None):
     try:
         if client is None:
             client = boto3.client("s3")
         if bucket_name is None:
-            bucket_name = transform_bucket(client)
+            bucket_name = get_transform_bucket()
         files = client.list_objects_v2(Bucket=bucket_name)
 
         dfs = {}
@@ -96,7 +103,7 @@ def convert_parquet_files_to_dfs(bucket_name=None, client=None):
                     logger.error(f"Unable to process file {file_key}: {e}")
         else:
             logger.error(f"No files found in {bucket_name}.")
-            return []
+            return {}
     except ValueError as value_error:
         logger.error(f"Unable to list objects: {value_error}")
         raise
@@ -111,11 +118,12 @@ def upload_dfs_to_database():
     dict_of_dfs = convert_parquet_files_to_dfs()
     db_engine = connect_to_db_and_return_engine()
     try:
-        for table_name, df in dict_of_dfs:
-            df.to_sql(table_name, con=db_engine, ifexists="replace", index=False)
+        for table_name, df in dict_of_dfs.items():
+            df.to_sql(table_name, con=db_engine, if_exists="replace", index=False)
             uploaded.append(table_name)
     except Exception as e:
         logger.error(f"Error uploading dataframes: {e}")
+        raise
     db_engine.dispose()
     return uploaded