Merge branch 'development' into test/tests_transform_lambda

14 files changed, 422 insertions, 191 deletions

diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml
new file mode 100644
index 0000000..5672048
--- /dev/null
+++ b/.github/workflows/deploy.yml
@@ -0,0 +1,42 @@
+name: deploy-terraform
+
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+      
+
+jobs:
+  deploy-terraform:
+    name: Deploy Terraform
+    runs-on: ubuntu-latest
+    #needs: run-checks (must ref on-commit.yml file)
+    environment: production
+    steps:
+      - name: Checkout Repo
+        uses: actions/checkout@v4
+      
+      - name: Install Terraform
+        uses: hashicorp/setup-terraform@v3
+
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          aws-region: ${{ secrets.AWS_REGION }}
+
+      - name: Terraform Init
+        working-directory: terraform
+        run: terraform init
+
+      - name: Terraform Plan
+        working-directory: terraform
+        run: terraform plan
+
+      - name: Terraform Apply
+        working-directory: terraform
+        run: terraform apply --auto-approve
diff --git a/.gitignore b/.gitignore
index 80f83ae..d5e8d2d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,10 +11,10 @@
 *.zip
 log*
 __pycache__/
+*.parquet
+/dim_*
+/fact_*
 
 # OS-Related Files
 .DS_Store
-venv
-
-##testing
-*.parquet
\ No newline at end of file
+venv
\ No newline at end of file
diff --git a/README.md b/README.md
index 7d7e499..e8e1ebd 100644
--- a/README.md
+++ b/README.md
@@ -36,16 +36,10 @@ can:
 We aim for the project to have certain features. Some are more prioritised than
 others.
 
-- [ ] Automated data ingestion from `totesys` db
-- [ ] Data storage for ingested and processed data in S3 buckets
-- [ ] Data transformation for data warehouse schema
-- [ ] Automated data loading into the data warehouse schema
-- [ ] Logging and monitoring with CloudWatch
-- [ ] Notifications for errors and successful runs (e.g. successful ingestion)
-- [ ] Visualisation of warehouse data
-
-# Test Coverage
-TBA
-
-# Contributors
-TBA
+- Automated data ingestion from `totesys` db
+- Data storage for ingested and processed data in S3 buckets
+- Data transformation for data warehouse schema
+- Automated data loading into the data warehouse schema
+- Logging and monitoring with CloudWatch
+- Notifications for errors and successful runs (e.g. successful ingestion)
+- Visualisation of warehouse data
diff --git a/car_data.parquet b/car_data.parquet
deleted file mode 100644
index 1853af6..0000000
--- a/car_data.parquet
+++ /dev/null
diff --git a/requirements_lambda_01.txt b/requirements_lambda_01.txt
new file mode 100644
index 0000000..10f56be
--- /dev/null
+++ b/requirements_lambda_01.txt
@@ -0,0 +1,3 @@
+boto3
+botocore
+pg8000
\ No newline at end of file
diff --git a/requirements_lambda_02.txt b/requirements_lambda_02.txt
new file mode 100644
index 0000000..20c88d7
--- /dev/null
+++ b/requirements_lambda_02.txt
@@ -0,0 +1,9 @@
+pandas
+pyarrow
+SQLAlchemy
+auto_mix_prep
+beautifulsoup4
+boto3
+botocore
+pg8000
+Requests
\ No newline at end of file
diff --git a/scripts/make_layer_zip.sh b/scripts/make_layer_zip.sh
index eabe301..7f64873 100755
--- a/scripts/make_layer_zip.sh
+++ b/scripts/make_layer_zip.sh
@@ -1,8 +1,17 @@
-# Description: Make the zip file for the layer
+# Description: Make the zip file for the layer for the extract lambda function
 
 cd "$(dirname "$0")/.."
+
+# Layer 01
+mkdir -p python/lib/python3.11/site-packages
+pip3 install --upgrade -r requirements_lambda_01.txt -t python/lib/python3.11/site-packages
+rm layer_01.zip
+zip -r layer_01.zip python
+rm -r python/
+
+# Layer 02
 mkdir -p python/lib/python3.11/site-packages
-pip3 install --upgrade -r requirements.txt -t python/lib/python3.11/site-packages
-rm layer.zip
-zip -r layer.zip python
+pip3 install --upgrade -r requirements_lambda_02.txt -t python/lib/python3.11/site-packages
+rm layer_02.zip
+zip -r layer_02.zip python
 rm -r python/
diff --git a/src/load_lambda.py b/src/load_lambda.py
index 7339ab9..86189dc 100644
--- a/src/load_lambda.py
+++ b/src/load_lambda.py
@@ -7,7 +7,8 @@ import logging
 import json
 import traceback
 from sqlalchemy import create_engine
-
+from datetime import datetime as dt
+import re
 
 logger = logging.getLogger(__name__)
 
@@ -15,10 +16,10 @@ logging.basicConfig(
     format="{asctime} - {levelname} - {message}",
     style="{",
     datefmt="%Y-%m-%d %H:%M",
-    level=logging.DEBUG,
+    level=logging.INFO,
 )
-
-logging.getLogger("botocore").setLevel(logging.INFO)
+# logging.getLogger("botocore").setLevel(logging.INFO)
+# logging.getLogger('sqlalchemy.engine').setLevel(logging.DEBUG)
 
 
 def lambda_handler(event, context):
@@ -38,10 +39,10 @@ def lambda_handler(event, context):
                 ),
             }
         else:
-            logger.error(f"error")
+            logger.error(f"error", exc_info=True)
             return {"error"}
     except Exception as e:
-        logger.error({e})
+        logger.error({e}, exc_info=True)
         return {"statusCode": 500, "body": {e}}
 
 
@@ -56,12 +57,15 @@ def retrieve_secrets(client=None, secret_name=None):
 
     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)}")
+        logger.error(
+            f"Failed to retrieve secret {secret_name}: {str(e)}", exc_info=True
+        )
         raise e
     except KeyError:
-        logger.error(f"Secret {secret_name} does not contain a SecretString")
+        logger.error(
+            f"Secret {secret_name} does not contain a SecretString", exc_info=True
+        )
         raise ValueError(f"Secret {secret_name} does not contain a SecretString")
 
     return get_secret_value_response["SecretString"]
@@ -86,7 +90,7 @@ def connect_to_db_and_return_engine(sm_secret=None):
         engine = create_engine(conn_str)
         return engine
     except Exception as e:
-        logger.error(f"Interface error: {e}")
+        logger.error(f"Interface error: {e}", exc_info=True)
         raise RuntimeError("Failed to create database engine")
 
 
@@ -97,7 +101,7 @@ def get_transform_bucket(client=None):
     try:
         response = client.list_buckets()
     except ClientError as e:
-        logger.error(f"Error listing S3 buckets: {e}")
+        logger.error(f"Error listing S3 buckets: {e}", exc_info=True)
         raise RuntimeError("Error listing S3 buckets")
 
     transform_bucket_filter = [
@@ -107,7 +111,7 @@ def get_transform_bucket(client=None):
     ]
 
     if not transform_bucket_filter:
-        logger.error("No transform bucket found")
+        logger.error("No transform bucket found", exc_info=True)
         raise ValueError("No transform bucket found")
 
     return transform_bucket_filter[0]
@@ -118,7 +122,26 @@ def get_transform_bucket(client=None):
 # return a dictionary of dataframes with name as key, and dataframe object as value
 
 
+def get_latest_timestamp(existing_files):
+    if existing_files:
+        all_datetimes = []
+        for file_name in existing_files:
+            match = re.search(r"\/(.+/).+_(.+)\.parquet", file_name)
+            if match:
+                datetime_str = "".join(match.group(1, 2))
+                all_datetimes.append(dt.strptime(datetime_str, "%Y/%m/%d/%H:%M:%S"))
+        return max(all_datetimes) if all_datetimes else dt.min
+    return existing_files
+
+
 def convert_parquet_files_to_dfs(bucket_name=None, client=None):
+    mutable_df_dict = [
+        "dim_currency",
+        "fact_sales_order",
+        "fact_purchase_order",
+        "fact_payment",
+    ]
+
     try:
         if client is None:
             client = boto3.client("s3")
@@ -128,27 +151,53 @@ def convert_parquet_files_to_dfs(bucket_name=None, client=None):
 
         dfs = {}
         if "Contents" in files:
-            for file in files["Contents"]:
-                file_key = file["Key"]
+            s3_key_list = [file["Key"] for file in files["Contents"]]
+            immutables_l = []
+            mutables_d = {prefix: [] for prefix in mutable_df_dict}
+            for tab, s3_key in mutables_d.items():
+                for file in s3_key_list:
+                    if tab in file:
+                        s3_key.append(file)
+                    elif "2024" not in file:
+                        immutables_l.append(file)
+                    else:
+                        continue
+            immutables_l = list(set(immutables_l))
+            latest_s3_keys = []
+            for k, v in mutables_d.items():
+                latest_s3_keys.append(
+                    dt.strftime(
+                        get_latest_timestamp(v), f"{k}/%Y/%m/%d/{k}_%H:%M:%S.parquet"
+                    )
+                )
+            for file_key in immutables_l + latest_s3_keys:
                 try:
                     file_obj = client.get_object(Bucket=bucket_name, Key=file_key)
                     parquet_file = pq.ParquetFile(BytesIO(file_obj["Body"].read()))
                     df = parquet_file.read().to_pandas()
-                    dfs[file_key] = df
+                    # >> can't do 'any' (default) because we lose rows in dim_location
+                    df_without_nulls = df.dropna(how="all")
+                    # print("df_without_nulls", df_without_nulls)
+                    # print("type", type(df_without_nulls))
+                    # print(df_without_nulls.columns)
+                    dfs[file_key] = df_without_nulls
                 except ClientError as e:
-                    logger.error(f"Unable to retrieve S3 object {file_key}: {e}")
+                    logger.error(
+                        f"Unable to retrieve S3 object {file_key}: {e}", exc_info=True
+                    )
                 except Exception as e:
-                    logger.error(f"Unable to process file {file_key}: {e}")
+                    logger.error(
+                        f"Unable to process file {file_key}: {e}", exc_info=True
+                    )
         else:
-            logger.error(f"No files found in {bucket_name}.")
+            logger.error(f"No files found in {bucket_name}.", exc_info=True)
             return {}
     except ValueError as value_error:
-        logger.error(f"Unable to list objects: {value_error}")
+        logger.error(f"Unable to list objects: {value_error}", exc_info=True)
         raise
     except ClientError as client_error:
-        logger.error(f"Unable to list objects: {client_error}")
+        logger.error(f"Unable to list objects: {client_error}", exc_info=True)
         raise
-
     return dfs
 
 
@@ -162,47 +211,65 @@ def upload_dfs_to_database():
         "dim_location.parquet",
         "dim_staff.parquet",
         "dim_design.parquet",
+        "dim_transaction.parquet",  # This one was missing,
+        "dim_payment_type.parquet",
     ]
     mutable_df_dict = [
+        "dim_currency",
         "fact_sales_order",
         "fact_purchase_order",
         "fact_payment",
-        "dim_currency",
     ]
-
-    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,
-                    if_exists="append",
-                    index=False,
-                )
-                upload_status["uploaded"].append(table_name)
-            except Exception as e:
-                logger.error(f"Error uploading dataframe {file_name} to database: {e}")
-                raise
-        elif file_name.rsplit("_", 1)[0] in mutable_df_dict:
-            table_name = file_name.rsplit("_", 1)[0]
-            try:
-                df.to_sql(
-                    table_name,
-                    con=db_engine,
-                    schema="project_team_2",
-                    if_exists="append",
-                    index=False,
+    with db_engine.begin() as connection:
+        for file_name, df in dict_of_dfs.items():
+            print(df.dtypes, "dtypes")
+            print(df.head())
+            print(file_name, "<<< FILE NAME")
+            print(immutable_df_dict, "<<<IMMUTABLE_DF_DICT")
+            if file_name in immutable_df_dict:
+                table_name = file_name.split(".")[0]
+                print(table_name, "<<<<<")
+                try:
+                    df.to_sql(
+                        table_name,
+                        con=connection,
+                        schema="project_team_2",
+                        if_exists="append",
+                        index=False,
+                    )
+                    upload_status["uploaded"].append(table_name)
+                    print(upload_status)
+                except Exception as e:
+                    logger.error(
+                        f"Error uploading dataframe {file_name} to database: {e}",
+                        exc_info=True,
+                    )
+                    raise
+            elif file_name.split("/")[0] in mutable_df_dict:
+                table_name = file_name.split("/")[0]
+                print(table_name, "<<<<<<<TABLE NAME")
+                try:
+                    df.to_sql(
+                        table_name,
+                        con=connection,
+                        schema="project_team_2",
+                        if_exists="append",
+                        index=False,
+                    )
+                    upload_status["uploaded"].append(table_name)
+                except Exception as e:
+                    logger.error(
+                        f"Error uploading dataframe {file_name} to database: {e}",
+                        exc_info=True,
+                    )
+                    raise
+            else:
+                upload_status["not_uploaded"].append(file_name)
+                logger.error(
+                    f"{file_name} does not correspond with table in database",
+                    exc_info=True,
                 )
-                upload_status["uploaded"].append(table_name)
-            except Exception as e:
-                logger.error(f"Error uploading dataframe {file_name} to database: {e}")
-                raise
-        else:
-            upload_status["not_uploaded"].append(file_name)
-            logger.error(f"{file_name} does not correspond with table in database")
+            print(upload_status)
     db_engine.dispose()
     return upload_status
 
diff --git a/src/transform_lambda/dataframes.py b/src/transform_lambda/dataframes.py
index 2a46bd6..6de58e7 100644
--- a/src/transform_lambda/dataframes.py
+++ b/src/transform_lambda/dataframes.py
@@ -18,8 +18,7 @@ import requests
 
 # no test, same as fact_payment
 def create_fact_sales_order(dict_of_df):
-    df_sales = dict_of_df["sales_order"]
-    df_sales.index.name = "sales_record_id"
+    df_sales = dict_of_df["sales_order"].rename(columns={"staff_id": "sales_staff_id"})
 
     df_sales["created_date"] = df_sales["created_at"].astype("datetime64[ns]").dt.date
     df_sales["created_time"] = (
@@ -37,30 +36,31 @@ def create_fact_sales_order(dict_of_df):
     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
-    )
-    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
+    fact_sales = df_sales.loc[
+        :,
+        [
+            "sales_order_id",
+            "created_date",
+            "created_time",
+            "last_updated_date",
+            "last_updated_time",
+            "sales_staff_id",
+            "counterparty_id",
+            "units_sold",
+            "unit_price",
+            "currency_id",
+            "design_id",
+            "agreed_payment_date",
+            "agreed_delivery_date",
+            "agreed_delivery_location_id",
+        ],
+    ]
+    fact_sales.convert_dtypes()
+    fact_sales.index = pd.RangeIndex(1, len(fact_sales.index) + 1)
+    fact_sales.index.name = "sales_record_id"
+    fact_sales.reset_index(inplace=True)
+    fact_sales.dropna(inplace=True)
+    return fact_sales
 
 
 # no test, same as fact_payment
@@ -68,7 +68,6 @@ def create_fact_sales_order(dict_of_df):
 
 def create_fact_purchase_orders(dict_of_df):
     df_po = dict_of_df["purchase_order"]
-    df_po.index.name = "purchase_record_id"
     df_po["created_date"] = df_po["created_at"].astype("datetime64[ns]").dt.date
     df_po["created_time"] = (
         df_po["created_at"].astype("datetime64[ns]").dt.floor("s").dt.time
@@ -83,9 +82,31 @@ def create_fact_purchase_orders(dict_of_df):
     df_po["agreed_payment_date"] = pd.to_datetime(
         df_po["agreed_payment_date"], format="%Y-%m-%d"
     )
-    df_po = df_po.drop(labels=["created_at", "last_updated"], axis=1)
-    df_po.reset_index(inplace=True)
-    return df_po
+    fact_purchase_order = df_po.loc[
+        :,
+        [
+            "purchase_order_id",
+            "created_date",
+            "created_time",
+            "last_updated_date",
+            "last_updated_time",
+            "staff_id",
+            "counterparty_id",
+            "item_code",
+            "item_quantity",
+            "item_unit_price",
+            "currency_id",
+            "agreed_delivery_date",
+            "agreed_payment_date",
+            "agreed_delivery_location_id",
+        ],
+    ]
+    fact_purchase_order.convert_dtypes()
+    fact_purchase_order.index = pd.RangeIndex(1, len(fact_purchase_order.index) + 1)
+    fact_purchase_order.index.name = "purchase_record_id"
+    fact_purchase_order.reset_index(inplace=True)
+    fact_purchase_order.dropna(inplace=True)
+    return fact_purchase_order
 
 
 # test passed
@@ -93,7 +114,6 @@ def create_fact_purchase_orders(dict_of_df):
 
 def create_fact_payment(dict_of_df):
     df_payment = dict_of_df["payment"]
-    df_payment.index.name = "payment_record_id"
     df_payment["created_date"] = (
         df_payment["created_at"].astype("datetime64[ns]").dt.date
     )
@@ -109,38 +129,60 @@ def create_fact_payment(dict_of_df):
     df_payment["payment_date"] = pd.to_datetime(
         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
+    fact_payment = df_payment.loc[
+        :,
+        [
+            "payment_id",
+            "created_date",
+            "created_time",
+            "last_updated_date",
+            "last_updated_time",
+            "transaction_id",
+            "counterparty_id",
+            "payment_amount",
+            "currency_id",
+            "payment_type_id",
+            "paid",
+            "payment_date",
+        ],
+    ]
+    fact_payment.convert_dtypes()
+    fact_payment.index = pd.RangeIndex(1, len(fact_payment.index) + 1)
+    fact_payment.index.name = "payment_record_id"
+    fact_payment.reset_index(inplace=True)
+    fact_payment.dropna(inplace=True)
+    fact_payment = fact_payment.astype({"currency_id": "int", "payment_id": "int"})
+    return fact_payment
 
 
 # test passed
 
 
 def create_dim_transaction(dict_of_df):
-    df_transaction = dict_of_df["transaction"].drop(
-        labels=["created_at", "last_updated"], axis=1
-    )
-    return df_transaction
+    dim_transaction = dict_of_df["transaction"].loc[
+        :, ["transaction_id", "transaction_type", "sales_order_id", "purchase_order_id"]
+    ]
+    # dim_transaction = dim_transaction.astype({"sales_order_id":"Int64","purchase_order_id":"Int64"})
+    return dim_transaction
 
 
 # test passed
 
 
 def create_dim_location(dict_of_df):
-    df_loc = (
+    dim_location = (
         dict_of_df["address"]
         .drop(labels=["created_at", "last_updated"], axis=1)
         .rename(columns={"address_id": "location_id"})
     )
-    return df_loc
+    return dim_location
 
 
 def create_dim_counterparty(dict_of_df):
     df_prefixed_address = (
         dict_of_df["address"]
         .drop(labels=["created_at", "last_updated"], axis=1)
+        .rename(columns={"phone": "phone_number"})
         .add_prefix("counterparty_legal_", axis=1)
     )
     df_cp = pd.merge(
@@ -149,15 +191,19 @@ def create_dim_counterparty(dict_of_df):
         left_on="legal_address_id",
         right_on="counterparty_legal_address_id",
         how="inner",
-    )
-    df_cp.drop(
-        columns=[
+    )  # .dropna(inplace=True)
+    dim_counterparty = df_cp.drop(
+        labels=[
             "legal_address_id",
             "counterparty_legal_address_id",
+            "created_at",
+            "last_updated",
+            "commercial_contact",
+            "delivery_contact",
         ],
-        inplace=True,
+        axis=1,
     )
-    return df_cp
+    return dim_counterparty
 
 
 # test passed
@@ -179,6 +225,7 @@ def create_dim_date(dict_of_df):
     sr_date = pd.array(pd.concat(list_of_date_columns), dtype="datetime64[ns]")
     df_date = pd.DataFrame(data=sr_date, columns=["date_id"])
     df_date.drop_duplicates(inplace=True)
+    # df_date.dropna(inplace=True)
     df_date["year"] = df_date["date_id"].dt.year
     df_date["month"] = df_date["date_id"].dt.month
     df_date["day"] = df_date["date_id"].dt.day
@@ -210,10 +257,13 @@ def scrape_currency_names():
 
 def create_dim_currency(dict_of_df, names=scrape_currency_names()):
     df_cur = dict_of_df["currency"].drop(labels=["created_at", "last_updated"], axis=1)
-    dim_cur = pd.merge(
-        df_cur, names, left_on="currency_code", right_on="currency_code", how="inner"
+    dim_currency = pd.merge(
+        df_cur, names, left_on="currency_code", right_on="currency_code", how="left"
     )
-    return dim_cur
+    dim_currency.drop_duplicates(inplace=True)
+    dim_currency.astype({"currency_name": "string", "currency_code": "string"})
+    print(dim_currency.dtypes, "<<<<<<<<<Dtype")
+    return dim_currency
 
 
 # tests passed
diff --git a/src/transform_lambda/transform_lambda.py b/src/transform_lambda/transform_lambda.py
index 478b257..54d7d48 100644
--- a/src/transform_lambda/transform_lambda.py
+++ b/src/transform_lambda/transform_lambda.py
@@ -65,6 +65,8 @@ def lambda_handler(event, context):
             "dim_location": create_dim_location(dict_of_df),
             "dim_staff": create_dim_staff(dict_of_df),
             "dim_design": create_dim_design(dict_of_df),
+            "dim_transaction": create_dim_transaction(dict_of_df),
+            "dim_payment_type": create_dim_payment_type(dict_of_df),
         }
 
         mutable_df_dict = {
@@ -73,7 +75,8 @@ def lambda_handler(event, context):
             "fact_payment": create_fact_payment(dict_of_df),
             "dim_currency": create_dim_currency(dict_of_df),
         }
-
+        print(immutable_df_dict.values())
+        print(mutable_df_dict.values())
         status = process_to_parquet_and_upload_to_s3(
             existing_s3_files, immutable_df_dict, mutable_df_dict, bucket
         )
diff --git a/terraform/events.tf b/terraform/events.tf
index 53ae10a..7f8f641 100644
--- a/terraform/events.tf
+++ b/terraform/events.tf
@@ -86,7 +86,7 @@ resource "aws_s3_bucket_notification" "extract_bucket_notification" {
 resource "aws_lambda_permission" "allow_s3_transform_bucket" {
   statement_id  = "AllowS3InvokeLambdaTransform${random_string.s3_transform_suffix.result}"
   action        = "lambda:InvokeFunction"
-  function_name = aws_lambda_function.transform_lambda.function_name
+  function_name = aws_lambda_function.load_lambda.function_name
   principal     = "s3.amazonaws.com"
   source_arn    = aws_s3_bucket.transform_bucket.arn
 
@@ -102,7 +102,7 @@ resource "aws_s3_bucket_notification" "transform_bucket_notification" {
 
   lambda_function {
     events              = ["s3:ObjectCreated:*"]
-    lambda_function_arn = aws_lambda_function.transform_lambda.arn
+    lambda_function_arn = aws_lambda_function.load_lambda.arn
   }
 
   depends_on = [aws_lambda_permission.allow_s3_transform_bucket]
diff --git a/terraform/lambda.tf b/terraform/lambda.tf
index 5f4a58e..1e12180 100644
--- a/terraform/lambda.tf
+++ b/terraform/lambda.tf
@@ -3,46 +3,65 @@
 ####################
 
 locals {
-  layer_dir      = "../"
-  layer_zip      = "layer.zip"
-  layer_name     = "lambda_layer"
-  script_dir     = "../scripts"
-  layer_zip_path = "${local.layer_dir}/${local.layer_zip}"
+  layer_dir         = "../"
+  layer_zip_01      = "layer_01.zip"
+  layer_zip_02      = "layer_02.zip"
+  layer_name_01     = "lambda_layer_01"
+  layer_name_02     = "lambda_layer_02"
+  script_dir        = "../scripts"
+  layer_zip_01_path = "${local.layer_dir}${local.layer_zip_01}"
+  layer_zip_02_path = "${local.layer_dir}${local.layer_zip_02}"
 }
 
-######################
-# Lambda Layer Setup #
-######################
-
 resource "null_resource" "prepare_layer" {
-
-  # New change: only run the script if the layer zip does not exist
-
+  provisioner "local-exec" {
+    command = "bash ${local.script_dir}/make_layer_zip.sh"
+  }
   triggers = {
-    layer_zip_exists = fileexists(local.layer_zip_path) ? "exists" : "not_exists"
+    always_run = timestamp()
   }
+}
 
-  provisioner "local-exec" {
-    command = "if [ ! -f ${local.layer_zip_path} ]; then bash ${local.script_dir}/make_layer_zip.sh; fi"
-  }
+################################
+# Lambda Layer (Extract) Setup #
+################################
+resource "aws_s3_object" "lambda_layer_zip_01" {
+  bucket        = aws_s3_bucket.lambda_code_bucket.id #bucket instead of id
+  key           = "${local.layer_name_01}/${local.layer_zip_01}"
+  source        = "${local.layer_dir}${local.layer_zip_01}"
+  depends_on    = [null_resource.prepare_layer]
+  etag          = fileexists(local.layer_zip_01_path) ? filemd5(local.layer_zip_01_path) : null
+  force_destroy = true
+}
+
+resource "aws_lambda_layer_version" "lambda_layer_01" {
+  layer_name          = local.layer_name_01
+  compatible_runtimes = ["python3.11"]
+  s3_bucket           = aws_s3_bucket.lambda_code_bucket.bucket
+  s3_key              = aws_s3_object.lambda_layer_zip_01.key
+  source_code_hash    = fileexists(local.layer_zip_01_path) ? filebase64sha256(local.layer_zip_01_path) : null
+  depends_on          = [aws_s3_object.lambda_layer_zip_01]
 }
 
-resource "aws_s3_object" "lambda_layer_zip" {
-  bucket     = aws_s3_bucket.lambda_code_bucket.id #bucket instead of id
-  key        = "${local.layer_name}/${local.layer_zip}"
-  source     = "${local.layer_dir}/${local.layer_zip}"
-  depends_on = [null_resource.prepare_layer]
-  etag       = fileexists(local.layer_zip_path) ? filemd5(local.layer_zip_path) : null
+#########################################
+# Lambda Layer (Load & Transform) Setup #
+#########################################
+resource "aws_s3_object" "lambda_layer_zip_02" {
+  bucket        = aws_s3_bucket.lambda_code_bucket.id #bucket instead of id
+  key           = "${local.layer_name_02}/${local.layer_zip_02}"
+  source        = "${local.layer_dir}${local.layer_zip_02}"
+  depends_on    = [null_resource.prepare_layer]
+  etag          = fileexists(local.layer_zip_02_path) ? filemd5(local.layer_zip_02_path) : null
+  force_destroy = true
 }
 
-resource "aws_lambda_layer_version" "lambda_layer" {
-  layer_name          = local.layer_name
+resource "aws_lambda_layer_version" "lambda_layer_02" {
+  layer_name          = local.layer_name_02
   compatible_runtimes = ["python3.11"]
   s3_bucket           = aws_s3_bucket.lambda_code_bucket.bucket
-  s3_key              = aws_s3_object.lambda_layer_zip.key
-  source_code_hash    = fileexists(local.layer_zip_path) ? filebase64sha256(local.layer_zip_path) : null
-  skip_destroy        = true
-  depends_on          = [aws_s3_object.lambda_layer_zip]
+  s3_key              = aws_s3_object.lambda_layer_zip_02.key
+  source_code_hash    = fileexists(local.layer_zip_02_path) ? filebase64sha256(local.layer_zip_02_path) : null
+  depends_on          = [aws_s3_object.lambda_layer_zip_02]
 }
 
 ###########################
@@ -65,7 +84,7 @@ resource "aws_lambda_function" "extract_lambda" {
   function_name    = var.extract_lambda_name
   s3_bucket        = aws_s3_bucket.lambda_code_bucket.bucket
   s3_key           = aws_s3_object.extract_lambda_code.key
-  layers           = [aws_lambda_layer_version.lambda_layer.arn]
+  layers           = [aws_lambda_layer_version.lambda_layer_01.arn]
   role             = aws_iam_role.multi_service_role.arn
   handler          = "extract_lambda.lambda_handler"
   runtime          = "python3.11"
@@ -86,8 +105,8 @@ resource "aws_lambda_function" "extract_lambda" {
 
 data "archive_file" "transform_lambda_zip" {
   type        = "zip"
-  source_dir  = "${path.module}../src/transform_lambda" 
-  output_path = "${path.module}../transform_lambda.zip"
+  source_dir  = "${path.module}/../src/transform_lambda"
+  output_path = "${path.module}/../transform_lambda.zip"
 }
 
 resource "aws_s3_object" "transform_lambda_code" {
@@ -101,7 +120,7 @@ resource "aws_lambda_function" "transform_lambda" {
   function_name    = var.transform_lambda_name
   s3_bucket        = aws_s3_bucket.lambda_code_bucket.bucket
   s3_key           = aws_s3_object.transform_lambda_code.key
-  layers           = [aws_lambda_layer_version.lambda_layer.arn]
+  layers           = [aws_lambda_layer_version.lambda_layer_02.arn]
   role             = aws_iam_role.multi_service_role.arn
   handler          = "transform_lambda.lambda_handler"
   runtime          = "python3.11"
@@ -135,7 +154,7 @@ resource "aws_lambda_function" "load_lambda" {
   function_name    = var.load_lambda_name
   s3_bucket        = aws_s3_bucket.lambda_code_bucket.bucket
   s3_key           = aws_s3_object.load_lambda_code.key
-  layers           = [aws_lambda_layer_version.lambda_layer.arn]
+  layers           = [aws_lambda_layer_version.lambda_layer_02.arn]
   role             = aws_iam_role.multi_service_role.arn
   handler          = "load_lambda.lambda_handler"
   runtime          = "python3.11"
diff --git a/tests/test_dataframes.py b/tests/test_dataframes.py
index ea7bad1..7dd592a 100644
--- a/tests/test_dataframes.py
+++ b/tests/test_dataframes.py
@@ -1,4 +1,4 @@
-from src.dataframes import *
+from src.transform_lambda.dataframes import *
 import pandas as pd
 from unittest.mock import patch
 from datetime import datetime as dt
diff --git a/tests/test_load_lambda.py b/tests/test_load_lambda.py
index 65106f7..b284588 100644
--- a/tests/test_load_lambda.py
+++ b/tests/test_load_lambda.py
@@ -1,13 +1,20 @@
 import pandas as pd
-import pyarrow.parquet as pq
-from io import BytesIO
 from moto import mock_aws
 import boto3
 import botocore.exceptions
 import os
 import pytest
-from src.load_lambda import *
+from src.load_lambda import (
+    lambda_handler,
+    retrieve_secrets,
+    connect_to_db_and_return_engine,
+    convert_parquet_files_to_dfs,
+    get_transform_bucket,
+    upload_dfs_to_database,
+)
 import tempfile
+import json
+from unittest.mock import MagicMock, patch
 
 
 @pytest.fixture(scope="class")
@@ -20,19 +27,20 @@ def aws_credentials():
 
 
 @pytest.fixture(scope="class")
-def mock_s3_client(aws_credentials):
+def mock_s3_client():
     with mock_aws():
         yield boto3.client("s3")
 
 
 @pytest.fixture(scope="class")
-def mock_sm_client(aws_credentials):
+def mock_sm_client():
     with mock_aws():
         yield boto3.client("secretsmanager")
 
 
 class TestLambdaHandler:
-    def test_lambda_handler_returns_200_and_table_name_if_uploaded(self, mocker):
+    @staticmethod
+    def test_lambda_handler_returns_200_and_table_name_if_uploaded(mocker):
         mocker.patch(
             "src.load_lambda.upload_dfs_to_database",
             return_value={"uploaded": ["table_one", "table_two"], "not_uploaded": []},
@@ -42,7 +50,8 @@ class TestLambdaHandler:
         assert "table_one" in result["body"]
         assert "table_two" in result["body"]
 
-    def test_lambda_handler_returns_200_and_table_name_if_not_uploaded(self, mocker):
+    @staticmethod
+    def test_lambda_handler_returns_200_and_table_name_if_not_uploaded(mocker):
         mocker.patch(
             "src.load_lambda.upload_dfs_to_database",
             return_value={"uploaded": [], "not_uploaded": ["table_one"]},
@@ -51,7 +60,8 @@ class TestLambdaHandler:
         assert result["statusCode"] == 200
         assert "No dataframes were uploaded" in result["body"]
 
-    def test_lambda_handler_returns_error_if_both_lists_empty(self, mocker):
+    @staticmethod
+    def test_lambda_handler_returns_error_if_both_lists_empty(mocker):
         mocker.patch(
             "src.load_lambda.upload_dfs_to_database",
             return_value={"uploaded": [], "not_uploaded": []},
@@ -63,7 +73,8 @@ class TestLambdaHandler:
 
 
 class TestRetrieveSecrets:
-    def test_retrieve_secrets_returns_dictionary(self, mock_sm_client):
+    @staticmethod
+    def test_retrieve_secrets_returns_dictionary(mock_sm_client):
         secret = {
             "cohort_id": "test_cohort_id",
             "user": "test_user_id",
@@ -81,7 +92,8 @@ class TestRetrieveSecrets:
 
         assert isinstance(result, dict)
 
-    def test_retrieve_secrets_returns_correct_keys_and_values(self, mock_sm_client):
+    @staticmethod
+    def test_retrieve_secrets_returns_correct_keys_and_values(mock_sm_client):
         secret_name = "test_secret"
 
         result = json.loads(retrieve_secrets(mock_sm_client, secret_name))
@@ -89,7 +101,8 @@ class TestRetrieveSecrets:
         assert result["user"] == "test_user_id"
         assert result["password"] == "test_password"
 
-    def test_retrieve_secrets_returns_client_error_if_no_secret(self, mock_sm_client):
+    @staticmethod
+    def test_retrieve_secrets_returns_client_error_if_no_secret(mock_sm_client):
         secret_name = "another_test_secret"
 
         with pytest.raises(botocore.exceptions.ClientError) as error:
@@ -97,7 +110,8 @@ class TestRetrieveSecrets:
 
 
 class TestConnectToDBAndReturnEngine:
-    def test_returns_unsuccessful_connection_when_wrong_credentials(self):
+    @staticmethod
+    def test_returns_unsuccessful_connection_when_wrong_credentials():
         sm_secret = {
             "host": "host",
             "port": "port",
@@ -111,11 +125,13 @@ class TestConnectToDBAndReturnEngine:
 
 
 class TestGetTransformBucket:
-    def test_raises_value_error_if_no_buckets(self, mock_s3_client):
+    @staticmethod
+    def test_raises_value_error_if_no_buckets(mock_s3_client):
         with pytest.raises(ValueError, match="No transform bucket found"):
             get_transform_bucket(mock_s3_client)
 
-    def test_raises_value_error_if_no_transform_bucket(self, mock_s3_client):
+    @staticmethod
+    def test_raises_value_error_if_no_transform_bucket(mock_s3_client):
         mock_s3_client.create_bucket(
             Bucket="extract_bucket",
             CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
@@ -123,7 +139,8 @@ class TestGetTransformBucket:
         with pytest.raises(ValueError, match="No transform bucket found"):
             get_transform_bucket(mock_s3_client)
 
-    def test_returns_transform_bucket_if_one_bucket(self, mock_s3_client):
+    @staticmethod
+    def test_returns_transform_bucket_if_one_bucket(mock_s3_client):
         mock_s3_client.create_bucket(
             Bucket="transform_bucket",
             CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
@@ -131,7 +148,8 @@ class TestGetTransformBucket:
         result = get_transform_bucket(mock_s3_client)
         assert result == "transform_bucket"
 
-    def test_only_returns_transform_bucket_if_several_buckets(self, mock_s3_client):
+    @staticmethod
+    def test_only_returns_transform_bucket_if_several_buckets(mock_s3_client):
         mock_s3_client.create_bucket(
             Bucket="another_test_bucket",
             CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
@@ -141,7 +159,8 @@ class TestGetTransformBucket:
 
 
 class TestConvertParquetToDfs:
-    def test_function_returns_empty_dictionary_if_no_files(self, mock_s3_client):
+    @staticmethod
+    def test_function_returns_empty_dictionary_if_no_files(mock_s3_client):
         mock_s3_client.create_bucket(
             Bucket="transform_bucket",
             CreateBucketConfiguration={"LocationConstraint": "eu-west-2"},
@@ -151,14 +170,8 @@ class TestConvertParquetToDfs:
         )
         assert result == {}
 
-    # def test_function_returns_dictionary_with_table_with_file_key():
-    #     # need to mock parquet file and upload to mock bucket
-    #     result = convert_parquet_files_to_dfs(bucket_name="transform_bucket", client=mock_s3_client)
-    #     assert "dim_staff" in result
-
-    def test_function_returns_dictionary_with_file_key_and_dataframe(
-        self, mock_s3_client
-    ):
+    @staticmethod
+    def test_function_returns_dictionary_with_file_key_and_dataframe(mock_s3_client):
         with tempfile.TemporaryDirectory() as tmp:
             d = {
                 "test": ["Hello", "Bye"],
@@ -190,7 +203,29 @@ class TestConvertParquetToDfs:
 
 
 class TestUploadDfsToDatabase:
-    # Full success test
-    # Partial success test
-    # Failure test
-    pass
+    @pytest.fixture
+    def mock_engine(self):
+        engine = MagicMock()
+        engine.dispose = MagicMock()
+        return engine
+
+    @pytest.fixture
+    def mock_df(self):
+        df = MagicMock(spec=pd.DataFrame)
+        df.to_sql = MagicMock()
+        return df
+
+    @staticmethod
+    def test_function_returns_dictionary_with_uploaded_and_not_uploaded_keys(
+        mock_engine, mock_df
+    ):
+        with patch(
+            "src.load_lambda.convert_parquet_files_to_dfs",
+            return_value={"dim_counterparty.parquet": mock_df},
+        ), patch(
+            "src.load_lambda.connect_to_db_and_return_engine", return_value=mock_engine
+        ):
+            result = upload_dfs_to_database()
+
+            assert "uploaded" in result
+            assert "not_uploaded" in result