From 45e025ac0c4ae8c721cb0b875fd0abd67cc2bc07 Mon Sep 17 00:00:00 2001
From: T-Aji <tolujbd2@gmail.com>
Date: Wed, 14 Aug 2024 15:53:11 +0100
Subject: test: passing test for function connect_to_database

---
 src/extract_lambda.py | 40 +++++++++++++++++++++++++---------------
 1 file changed, 25 insertions(+), 15 deletions(-)

(limited to 'src')

diff --git a/src/extract_lambda.py b/src/extract_lambda.py
index dc70590..6e94bba 100644
--- a/src/extract_lambda.py
+++ b/src/extract_lambda.py
@@ -1,6 +1,5 @@
 from pg8000.native import Connection, DatabaseError, InterfaceError
-from dotenv import load_dotenv
-import os
+from dotenv import dotenv_values
 import boto3
 import csv
 from botocore.exceptions import ClientError
@@ -9,16 +8,15 @@ import json
 
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
-load_dotenv()
-
-
-database = os.getenv('database')
-user = os.getenv('user')
-password = os.getenv('password')
-host = os.getenv('host')
-port = os.getenv('port')
 
+class DBConnectionException(Exception):
+    """Wraps pg8000.native Error or DatabaseError."""
 
+    def __init__(self, e):
+        """Initialise with provided error message."""
+        self.message = str(e)
+        super().__init__(self.message)
+        
 def lambda_handler(event, context):
     """This lambda function connects to the Totesys database, lists the contents of the ingestion bucket,
        and converts all tables to CSV and if any of those tables do not exist in, or are different to the ones in s3, it uploads them
@@ -53,8 +51,19 @@ def lambda_handler(event, context):
         if db:
             db.close()
 
-def connect_to_database():
+def get_config(path: str = ".env") -> dict:
+    return dotenv_values(path)
+
+
+def connect_to_database() -> Connection:
     try:
+        config = get_config()
+        host = config["host"]
+        port = config["port"]
+        user = config["user"]
+        password = config["password"]
+        database = config["database"]
+
         return Connection(
             database=database,
             user=user,
@@ -62,12 +71,13 @@ def connect_to_database():
             host=host,
             port=port
         )
-    except DatabaseError as e:
-        logger.error(f'Database error: {e}')
-        raise
+    # except DatabaseError as e:
+    #     logger.error(f'Database error: {e}')
+    #     raise 
     except InterfaceError as i:
         logger.error(f'Interface error: {i}')
-        raise
+        raise DBConnectionException("Failed to connect to database")
+
 
 
 def list_existing_s3_files(bucket_name='extract_bucket', client=boto3.client('s3')):
-- 
cgit v1.2.3