PgDB: add exit_on_error parameter to connect method

This commit is contained in:
Benjamin Renard 2021-11-23 13:08:24 +01:00
parent 8e4864f2d5
commit 31b994ad3c

View file

@ -24,6 +24,18 @@ class PgDBException(Exception):
super().__init__(error.format(*args, **kwargs)) 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): class PgDBDuplicatedSQLParameter(PgDBException, KeyError):
""" """
Raised when trying to set a SQL query parameter Raised when trying to set a SQL query parameter
@ -77,7 +89,7 @@ class PgDB:
self._conn = None self._conn = None
self.just_try = just_try self.just_try = just_try
def connect(self): def connect(self, exit_on_error=True):
""" Connect to PostgreSQL server """ """ Connect to PostgreSQL server """
if self._conn is None: if self._conn is None:
try: try:
@ -90,12 +102,15 @@ class PgDB:
host=self._host, host=self._host,
password=self._pwd password=self._pwd
) )
except Exception: except Exception as err:
log.fatal( log.fatal(
'An error occured during Postgresql database connection (%s@%s, database=%s).', 'An error occured during Postgresql database connection (%s@%s, database=%s).',
self._user, self._host, self._db, exc_info=1 self._user, self._host, self._db, exc_info=1
) )
if exit_on_error:
sys.exit(1) sys.exit(1)
else:
raise PgDBFailToConnect(self._host, self._user, self._db) from err
return True return True
def close(self): def close(self):