diff options
| -rw-r--r-- | src/dataframes.py | 123 | ||||
| -rw-r--r-- | src/load_lambda.py | 56 | ||||
| -rw-r--r-- | tests/test_dataframes.py | 34 | ||||
| -rw-r--r-- | tests/test_load_lambda.py | 126 |
4 files changed, 251 insertions, 88 deletions
diff --git a/src/dataframes.py b/src/dataframes.py index f122368..ab32fff 100644 --- a/src/dataframes.py +++ b/src/dataframes.py @@ -20,6 +20,29 @@ 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 @@ -45,23 +68,25 @@ 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=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 ) - df_po["last_updated_date"] = df_po["last_updated"].astype("datetime64[ns]").dt.date - df_po["last_updated_time"] = ( + df_po["last_updated_date"]=df_po["last_updated"].astype( + "datetime64[ns]").dt.date + df_po["last_updated_time"]=( df_po["last_updated"].astype("datetime64[ns]").dt.floor("s").dt.time + ) - df_po["agreed_delivery_date"] = pd.to_datetime( + df_po["agreed_delivery_date"]=pd.to_datetime( df_po["agreed_delivery_date"], format="%Y-%m-%d" ) - df_po["agreed_payment_date"] = pd.to_datetime( + 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=df_po.drop(labels=["created_at", "last_updated"], axis=1) df_po.reset_index(inplace=True) return df_po @@ -70,24 +95,26 @@ 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=dict_of_df["payment"] + df_payment.index.name="payment_record_id" + df_payment["created_date"]=( df_payment["created_at"].astype("datetime64[ns]").dt.date ) - df_payment["created_time"] = ( + df_payment["created_time"]=( df_payment["created_at"].astype("datetime64[ns]").dt.floor("s").dt.time ) - df_payment["last_updated_date"] = ( + df_payment["last_updated_date"]=( df_payment["last_updated"].astype("datetime64[ns]").dt.date ) - df_payment["last_updated_time"] = ( - df_payment["last_updated"].astype("datetime64[ns]").dt.floor("s").dt.time + df_payment["last_updated_time"]=( + df_payment["last_updated"].astype( + "datetime64[ns]").dt.floor("s").dt.time ) - df_payment["payment_date"] = pd.to_datetime( + 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=df_payment.drop(labels=["created_at", "last_updated"], axis=1) + df_payment.reset_index(inplace=True) return df_payment @@ -96,7 +123,7 @@ def create_fact_payment(dict_of_df): def create_dim_transaction(dict_of_df): - df_transaction = dict_of_df["transaction"].drop( + df_transaction=dict_of_df["transaction"].drop( labels=["created_at", "last_updated"], axis=1 ) return df_transaction @@ -106,7 +133,7 @@ def create_dim_transaction(dict_of_df): def create_dim_location(dict_of_df): - df_loc = ( + df_loc=( dict_of_df["address"] .drop(labels=["created_at", "last_updated"], axis=1) .rename(columns={"address_id": "location_id"}) @@ -115,18 +142,18 @@ def create_dim_location(dict_of_df): def create_dim_counterparty(dict_of_df): - df_prefixed_address = dict_of_df["address"].add_prefix( + 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( + 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 @@ -135,7 +162,7 @@ def create_dim_counterparty(dict_of_df): def create_dim_date(dict_of_df): - fact_dfs = [ + fact_dfs=[ create_fact_payment(dict_of_df), create_fact_purchase_orders(dict_of_df), create_fact_sales_order(dict_of_df), @@ -147,16 +174,16 @@ def create_dim_date(dict_of_df): ] for col in date_col_names: list_of_date_columns.append(df[col]) - sr_date = pd.array(pd.concat(list_of_date_columns), dtype="datetime64[ns]") - df_date = pd.DataFrame(data=sr_date, columns=["date_id"]) + 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["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 - df_date["day_of_week"] = df_date["date_id"].dt.dayofweek - df_date["day_name"] = df_date["date_id"].dt.day_name() - df_date["month_name"] = df_date["date_id"].dt.month_name() - df_date["quarter"] = df_date["date_id"].dt.quarter + 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 + df_date["day_of_week"]=df_date["date_id"].dt.dayofweek + df_date["day_name"]=df_date["date_id"].dt.day_name() + df_date["month_name"]=df_date["date_id"].dt.month_name() + df_date["quarter"]=df_date["date_id"].dt.quarter return df_date @@ -164,13 +191,13 @@ def create_dim_date(dict_of_df): def scrape_currency_names(): - response = requests.get("https://www.xe.com/currency/").content - soup = BeautifulSoup(response, "html.parser") - currency = [ + response=requests.get("https://www.xe.com/currency/").content + soup=BeautifulSoup(response, "html.parser") + currency=[ item.text for item in soup.findAll("a", attrs={"class": "sc-299dec64-6 fZPTSw"}) ] - sr = pd.Series(currency) - df_cur = sr.str.split(pat=" - ", expand=True).rename( + sr=pd.Series(currency) + df_cur=sr.str.split(pat=" - ", expand=True).rename( {0: "currency_code", 1: "currency_name"}, axis=1 ) return df_cur @@ -180,8 +207,9 @@ 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=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" ) return dim_cur @@ -191,8 +219,9 @@ def create_dim_currency(dict_of_df, names=scrape_currency_names()): def create_dim_payment_type(dict_of_df): - df_payment_type = dict_of_df["payment_type"] - dim_payment_type = df_payment_type.loc[:, ["payment_type_id", "payment_type_name"]] + df_payment_type=dict_of_df["payment_type"] + dim_payment_type=df_payment_type.loc[:, [ + "payment_type_id", "payment_type_name"]] return dim_payment_type @@ -200,8 +229,8 @@ def create_dim_payment_type(dict_of_df): def create_dim_design(dict_of_df): - df_design = dict_of_df["design"] - dim_design = df_design.loc[ + df_design=dict_of_df["design"] + dim_design=df_design.loc[ :, ["design_id", "design_name", "file_name", "file_location"] ] return dim_design @@ -211,10 +240,10 @@ def create_dim_design(dict_of_df): def create_dim_staff(dict_of_df): - staff_department = pd.merge( + staff_department=pd.merge( dict_of_df["staff"], dict_of_df["department"], on="department_id", how="left" ) - dim_staff = staff_department.loc[ + dim_staff=staff_department.loc[ :, [ "staff_id", 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) diff --git a/tests/test_dataframes.py b/tests/test_dataframes.py index c9ff43f..ff282eb 100644 --- a/tests/test_dataframes.py +++ b/tests/test_dataframes.py @@ -54,7 +54,8 @@ class TestCreateDimStaff: "email_address": ["Hello", "Bye"], "department_id": ["Hello", "Bye"], } - test_df = {"staff": pd.DataFrame(data=d), "department": pd.DataFrame(data=d2)} + test_df = {"staff": pd.DataFrame( + data=d), "department": pd.DataFrame(data=d2)} result = create_dim_staff(test_df) assert isinstance(result, pd.DataFrame) @@ -71,7 +72,10 @@ class TestCreateDimStaff: "email_address": ["Hello", "Bye"], "department_id": ["Hello", "Bye"], } - test_df = {"staff": pd.DataFrame(data=d), "department": pd.DataFrame(data=d2)} + + test_df = {"staff": pd.DataFrame( + data=d), "department": pd.DataFrame(data=d2)} + result = create_dim_staff(test_df) expected_d = { "staff_id": ["Hello", "Bye"], @@ -88,7 +92,9 @@ class TestCreateDimStaff: class TestCreatePaymentType: def test_create_dim_payment_type_returns_correct_columns_and_values(self): - d = {"payment_type_id": ["Hello", "Bye"], "payment_type_name": ["Hello", "Bye"]} + d = {"payment_type_id": ["Hello", "Bye"], + "payment_type_name": ["Hello", "Bye"]} + test_df = {"payment_type": pd.DataFrame(data=d)} result = create_dim_payment_type(test_df) expected_columns = ["payment_type_id", "payment_type_name"] @@ -180,11 +186,14 @@ class TestCreateDimDate: index=[0], ) df_two = pd.DataFrame( - data={"updated_date": dt(2020, 5, 17), "created_date": dt(2021, 9, 13)}, + data={"updated_date": dt(2020, 5, 17), + "created_date": dt(2021, 9, 13)}, index=[0], ) df_three = pd.DataFrame( - data={"updated_date": dt(2022, 5, 17), "created_date": dt(2023, 5, 13)}, + data={"updated_date": dt(2022, 5, 17), + "created_date": dt(2023, 5, 13)}, + index=[0], ) expected_df = pd.DataFrame( @@ -214,7 +223,9 @@ class TestCreateDimDate: mock_fso.return_value = df_three result = create_dim_date({"dum": 0}) result.reset_index(inplace=True, drop=True) - assert result.eq(expected_df, axis="columns").all(axis=None) + assert result.eq( + expected_df, axis="columns").all(axis=None) + class TestCreateDimLocation: @@ -222,7 +233,9 @@ class TestCreateDimLocation: dict_df = { "address": pd.DataFrame( data=[["some_time", "some_other_time", 1, "SE18 9QO"]], - columns=["created_at", "last_updated", "address_id", "postal_code"], + columns=["created_at", "last_updated", + "address_id", "postal_code"], + ) } result = create_dim_location(dict_df) @@ -287,5 +300,8 @@ class TestCreateFactPayment: for col in list(result.columns): assert col in expected_cols for col in expected_cols: - if "_date" or "_time" in col: - assert result[col].dtype == "O" + + + +if "_date" or "_time" in col: + assert result[col].dtype == "O" diff --git a/tests/test_load_lambda.py b/tests/test_load_lambda.py index 88c71e4..65106f7 100644 --- a/tests/test_load_lambda.py +++ b/tests/test_load_lambda.py @@ -3,22 +3,18 @@ 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 ( - lambda_handler, - connect_to_db_and_return_engine, - get_transform_bucket, - convert_parquet_files_to_dfs, - upload_dfs_to_database, -) +from src.load_lambda import * +import tempfile @pytest.fixture(scope="class") def aws_credentials(): os.environ["AWS_ACCESS_KEY_ID"] = "testing" os.environ["AWS_SECRET_ACCESS_KEY"] = "testing" - os.environ["AWS_SECURIT_TOKEN"] = "testing" + os.environ["AWS_SECURITY_TOKEN"] = "testing" os.environ["AWS_SESSION_TOKEN"] = "testing" os.environ["AWS_DEFAULT_REGION"] = "eu-west-2" @@ -29,16 +25,89 @@ def mock_s3_client(aws_credentials): yield boto3.client("s3") +@pytest.fixture(scope="class") +def mock_sm_client(aws_credentials): + with mock_aws(): + yield boto3.client("secretsmanager") + + class TestLambdaHandler: - pass + def test_lambda_handler_returns_200_and_table_name_if_uploaded(self, mocker): + mocker.patch( + "src.load_lambda.upload_dfs_to_database", + return_value={"uploaded": ["table_one", "table_two"], "not_uploaded": []}, + ) + result = lambda_handler(None, None) + assert result["statusCode"] == 200 + 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): + mocker.patch( + "src.load_lambda.upload_dfs_to_database", + return_value={"uploaded": [], "not_uploaded": ["table_one"]}, + ) + result = lambda_handler(None, None) + assert result["statusCode"] == 200 + assert "No dataframes were uploaded" in result["body"] + + def test_lambda_handler_returns_error_if_both_lists_empty(self, mocker): + mocker.patch( + "src.load_lambda.upload_dfs_to_database", + return_value={"uploaded": [], "not_uploaded": []}, + ) + + result = lambda_handler(None, None) + + assert result == {"error"} class TestRetrieveSecrets: - pass + def test_retrieve_secrets_returns_dictionary(self, mock_sm_client): + secret = { + "cohort_id": "test_cohort_id", + "user": "test_user_id", + "password": "test_password", + "host": "test_host", + "database": "test_database", + "port": "test_port", + } + + secret_name = "test_secret" + + mock_sm_client.create_secret(Name=secret_name, SecretString=json.dumps(secret)) + + result = json.loads(retrieve_secrets(mock_sm_client, secret_name)) + + assert isinstance(result, dict) + + def test_retrieve_secrets_returns_correct_keys_and_values(self, mock_sm_client): + secret_name = "test_secret" + + result = json.loads(retrieve_secrets(mock_sm_client, secret_name)) + + 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): + secret_name = "another_test_secret" + + with pytest.raises(botocore.exceptions.ClientError) as error: + retrieve_secrets(mock_sm_client, secret_name) class TestConnectToDBAndReturnEngine: - pass + def test_returns_unsuccessful_connection_when_wrong_credentials(self): + sm_secret = { + "host": "host", + "port": "port", + "user": "user", + "password": "password", + "database": "database", + } + + with pytest.raises(Exception): + connect_to_db_and_return_engine(json.dumps(sm_secret)) class TestGetTransformBucket: @@ -87,6 +156,41 @@ class TestConvertParquetToDfs: # 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 + ): + with tempfile.TemporaryDirectory() as tmp: + d = { + "test": ["Hello", "Bye"], + "design_id": ["Hello", "Bye"], + "design_name": ["Hello", "Bye"], + "file_name": ["Hello", "Bye"], + "file_location": ["Hello", "Bye"], + "Hello": ["Hello", "Bye"], + } + + test_df = pd.DataFrame(data=d) + + path = os.path.join(tmp, "test_parquet.parquet") + + test_df.to_parquet(path, engine="pyarrow") + + with open(path, "rb") as p: + mock_s3_client.put_object( + Bucket="transform_bucket", Key="test_parquet.parquet", Body=p.read() + ) + + result = convert_parquet_files_to_dfs( + bucket_name="transform_bucket", client=mock_s3_client + ) + + assert "test_parquet.parquet" in result + + pd.testing.assert_frame_equal(result["test_parquet.parquet"], test_df) + class TestUploadDfsToDatabase: + # Full success test + # Partial success test + # Failure test pass |