From 31b994ad3c301d8eacfd489fc1d6322a331c4e40 Mon Sep 17 00:00:00 2001
From: Benjamin Renard <brenard@easter-eggs.com>
Date: Tue, 23 Nov 2021 13:08:24 +0100
Subject: [PATCH] PgDB: add exit_on_error parameter to connect method

---
 mylib/pgsql.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/mylib/pgsql.py b/mylib/pgsql.py
index f2fb2d0..76d0773 100644
--- a/mylib/pgsql.py
+++ b/mylib/pgsql.py
@@ -24,6 +24,18 @@ class PgDBException(Exception):
         super().__init__(error.format(*args, **kwargs))
 
 
+class PgDBFailToConnect(PgDBException, RuntimeError):
+    """
+    Raised on connecting error occurred
+    """
+
+    def __init__(self, host, user, db):
+        super().__init__(
+            "An error occured during Postgresql database connection ({user}@{host}, database={db})",
+            user=user, host=host, db=db
+        )
+
+
 class PgDBDuplicatedSQLParameter(PgDBException, KeyError):
     """
     Raised when trying to set a SQL query parameter
@@ -77,7 +89,7 @@ class PgDB:
         self._conn = None
         self.just_try = just_try
 
-    def connect(self):
+    def connect(self, exit_on_error=True):
         """ Connect to PostgreSQL server """
         if self._conn is None:
             try:
@@ -90,12 +102,15 @@ class PgDB:
                     host=self._host,
                     password=self._pwd
                 )
-            except Exception:
+            except Exception as err:
                 log.fatal(
                     'An error occured during Postgresql database connection (%s@%s, database=%s).',
                     self._user, self._host, self._db, exc_info=1
                 )
-                sys.exit(1)
+                if exit_on_error:
+                    sys.exit(1)
+                else:
+                    raise PgDBFailToConnect(self._host, self._user, self._db) from err
         return True
 
     def close(self):