Merge pull request #102 from ajschofield/feature/load-lambda-tests

pr: feature/load lambda tests

2 files changed, 68 insertions, 25 deletions

diff --git a/src/dataframes.py b/src/dataframes.py
index f122368..2a46bd6 100644
--- a/src/dataframes.py
+++ b/src/dataframes.py
@@ -20,6 +20,28 @@ import requests
 def create_fact_sales_order(dict_of_df):
     df_sales = dict_of_df["sales_order"]
     df_sales.index.name = "sales_record_id"
+
+    df_sales["created_date"] = df_sales["created_at"].astype("datetime64[ns]").dt.date
+    df_sales["created_time"] = (
+        df_sales["created_at"].astype("datetime64[ns]").dt.floor("s").dt.time
+    )
+    df_sales["last_updated_date"] = (
+        df_sales["last_updated"].astype("datetime64[ns]").dt.date
+    )
+    df_sales["last_updated_time"] = (
+        df_sales["last_updated"].astype("datetime64[ns]").dt.floor("s").dt.time
+    )
+    df_sales["agreed_delivery_date"] = pd.to_datetime(
+        df_sales["agreed_delivery_date"], format="%Y-%m-%d"
+    )
+    df_sales["agreed_payment_date"] = pd.to_datetime(
+        df_sales["agreed_payment_date"], format="%Y-%m-%d"
+    )
+    df_sales = df_sales.drop(labels=["created_at", "last_updated"], axis=1)
+
+    df_sales.reset_index(inplace=True)
+    return df_sales
+
     df_sales["created_date"] = df_sales["created_at"].astype("datetime64[ns]").dt.date
     df_sales["created_time"] = (
         df_sales["created_at"].astype("datetime64[ns]").dt.floor("s").dt.time
@@ -88,6 +110,7 @@ def create_fact_payment(dict_of_df):
         df_payment["payment_date"], format="%Y-%m-%d"
     )
     df_payment = df_payment.drop(labels=["created_at", "last_updated"], axis=1)
+
     df_payment.reset_index(inplace=True)
     return df_payment
 
@@ -115,18 +138,24 @@ def create_dim_location(dict_of_df):
 
 
 def create_dim_counterparty(dict_of_df):
-    df_prefixed_address = dict_of_df["address"].add_prefix(
-        "counterparty_legal_", axis=1
+    df_prefixed_address = (
+        dict_of_df["address"]
+        .drop(labels=["created_at", "last_updated"], axis=1)
+        .add_prefix("counterparty_legal_", axis=1)
     )
     df_cp = pd.merge(
         dict_of_df["counterparty"],
         df_prefixed_address,
         left_on="legal_address_id",
         right_on="counterparty_legal_address_id",
-        how="outer",
+        how="inner",
     )
     df_cp.drop(
-        columns=["legal_address_id", "counterparty_legal_address_id"], inplace=True
+        columns=[
+            "legal_address_id",
+            "counterparty_legal_address_id",
+        ],
+        inplace=True,
     )
     return df_cp
 
diff --git a/src/load_lambda.py b/src/load_lambda.py
index 6e6bc80..7339ab9 100644
--- a/src/load_lambda.py
+++ b/src/load_lambda.py
@@ -5,7 +5,7 @@ import pyarrow.parquet as pq
 from io import BytesIO
 import logging
 import json
-from src.extract_lambda import retrieve_secrets
+import traceback
 from sqlalchemy import create_engine
 
 
@@ -24,33 +24,39 @@ logging.getLogger("botocore").setLevel(logging.INFO)
 def lambda_handler(event, context):
     try:
         uploaded_tables = upload_dfs_to_database()
-        if not uploaded_tables["uploaded"]:
+        if uploaded_tables["not_uploaded"]:
             return {
                 "statusCode": 200,
                 "body": json.dumps("No dataframes were uploaded."),
             }
-        return {
-            "statusCode": 200,
-            "body": json.dumps(
-                f"""The following dataframes were uploaded successfully: 
-                {uploaded_tables["uploaded"]} ."""
-            ),
-        }
+        elif uploaded_tables["uploaded"]:
+            return {
+                "statusCode": 200,
+                "body": json.dumps(
+                    f"""The following dataframes were uploaded successfully: 
+                    {uploaded_tables["uploaded"]} ."""
+                ),
+            }
+        else:
+            logger.error(f"error")
+            return {"error"}
     except Exception as e:
-        logger.error(f"Error: {e}", exc_info=True)
-        return {"statusCode": 500, "body": json.dumps("Internal server error.")}
+        logger.error({e})
+        return {"statusCode": 500, "body": {e}}
 
 
-def retrieve_secrets():
-    secret_name = "bentley-RDS-credentials"
+def retrieve_secrets(client=None, secret_name=None):
+    session = boto3.session.Session()
     region_name = "eu-west-2"
 
-    # Create a Secrets Manager client
-    session = boto3.session.Session()
-    client = session.client(service_name="secretsmanager", region_name=region_name)
+    if secret_name == None:
+        secret_name = "bentley-RDS-credentials"
+    if client == None:
+        client = session.client(service_name="secretsmanager", region_name=region_name)
 
     try:
         get_secret_value_response = client.get_secret_value(SecretId=secret_name)
+        print(get_secret_value_response)
     except ClientError as e:
         logger.error(f"Failed to retrieve secret {secret_name}: {str(e)}")
         raise e
@@ -64,9 +70,12 @@ def retrieve_secrets():
 # connect to database, slightly different way of doing it, to allow manipulation through pandas
 
 
-def connect_to_db_and_return_engine():
+def connect_to_db_and_return_engine(sm_secret=None):
+    if sm_secret is None:
+        sm_secret = json.loads(retrieve_secrets())
+
     try:
-        secrets = json.loads(retrieve_secrets())
+        secrets = sm_secret
         host = secrets["host"]
         port = secrets["port"]
         user = secrets["user"]
@@ -162,14 +171,15 @@ def upload_dfs_to_database():
     ]
 
     for file_name, df in dict_of_dfs.items():
+        print(df)
         if file_name in immutable_df_dict:
             table_name = file_name.split(".")[0]
+            print(table_name, "<<<<<")
             try:
                 df.to_sql(
                     table_name,
                     con=db_engine,
-                    schema="project_team_2",
-                    if_exists="overwrite",
+                    if_exists="append",
                     index=False,
                 )
                 upload_status["uploaded"].append(table_name)
@@ -183,7 +193,7 @@ def upload_dfs_to_database():
                     table_name,
                     con=db_engine,
                     schema="project_team_2",
-                    if_exists="overwrite",
+                    if_exists="append",
                     index=False,
                 )
                 upload_status["uploaded"].append(table_name)
@@ -195,3 +205,7 @@ def upload_dfs_to_database():
             logger.error(f"{file_name} does not correspond with table in database")
     db_engine.dispose()
     return upload_status
+
+
+if __name__ == "__main__":
+    lambda_handler(None, None)