upgrade the error handling at the db connection

Author: cortze

unit_db_test/postgresql.py

@@ -1,8 +1,16 @@
 # postgresql.py
 import pandas as pd
+import psycopg2 as pspg
 import sqlalchemy as sa
 
 
+class  NoConnectionToDBError(Exception):
+    def __init__(self, cause):
+        self.cause = cause
+
+    def cause(self):
+        return self.cause
+
 def parse_db_endpoint_string(s):
     user = s.split(":")[1].split("@")[0].split("//")[1]
     password = s.split(":")[2].split("@")[0]
@@ -15,23 +23,27 @@ def parse_db_endpoint_string(s):
 class Postgres:
     def __init__(
         self,
-        port=None,
-        user=None,
-        password=None,
-        database=None,
-        host="localhost",
-        url=None,
+        user: str,
+        password: str,
+        host: str,
+        port: str,
+        database: str,
     ):
-        self.port = port
         self.user = user
         self.password = password
+        self.host = host
+        self.port = port
         self.database = database
-        if url is not None:
-            self.engine = sa.create_engine(url=url)
-        else:
-            self.url = f"postgresql://{user}:{password}@{host}:{port}/{database}"
-            self.engine = sa.create_engine(url=url)
-        self.conn = self.engine.connect()
+        self.url = f"postgresql://{user}:{password}@{host}:{port}/{database}"
+        try:
+            self.engine = sa.create_engine(url=self.url)
+            self.conn = self.engine.connect()
+        except pspg.OperationalError as dbError:
+            print(f"Error connecting to the database: {dbError}")
+            raise NoConnectionToDBError(dbError)
+        except Exception as e:
+            print(f"An unexpected error occurred: {e}")
+            raise NoConnectionToDBError(e)
 
     def get_df_from_sql_query(self, sql_query):
         return pd.read_sql_query(sa.text(sql_query), self.conn)

unit_db_test/testcase.py

@@ -2,8 +2,8 @@
 import pandas as pd
 import unittest
 
-from unit_db_test.postgresql import Postgres
-from unit_db_test.utils import load_credentials_from
+from unit_db_test.postgresql import Postgres, NoConnectionToDBError
+from unit_db_test.utils import load_credentials_from, NoEnvVariableError
 
 
 
@@ -26,9 +26,10 @@ def setUp(self):
         try:
             u, p, i, port, db, url = self.readCredentialsFromEnv()
             self.connectDB(u, p, i, port, db, url)
-        except Exception as e:
-            self.assertRaise(e)
-            self.fail("unable to make connection to the DB")
+        except NoEnvVariableError as envError:
+            raise envError
+        except NoConnectionToDBError as dbError:
+            raise dbError
 
     def setDown(self):
         self.db.close()
@@ -47,7 +48,7 @@ def assertCustomNullItemsInColumn(self, df: pd.DataFrame, column: str, target: i
     def assertNoRows(self, df: pd.DataFrame):
         self.assertNRows(df, 0)
 
-    def assertNRows(self, df: pd.DataFrame, target_rows: str):
+    def assertNRows(self, df: pd.DataFrame, target_rows: int):
         if len(df) != target_rows:
-            raise AssertionError(f"df has {len(df)} items, expected {target_rows}")
+            raise AssertionError(f"df has {len(df)} rows, expected {target_rows}")