diff --git a/api/chalicelib/utils/pg_client.py b/api/chalicelib/utils/pg_client.py index 69a5b5a8b..4cfd8b0e3 100644 --- a/api/chalicelib/utils/pg_client.py +++ b/api/chalicelib/utils/pg_client.py @@ -111,6 +111,8 @@ class PostgresClient: def __enter__(self): if self.cursor is None: self.cursor = self.connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) + self.cursor.cursor_execute = self.cursor.execute + self.cursor.execute = self.__execute self.cursor.recreate = self.recreate_cursor return self.cursor @@ -136,6 +138,17 @@ class PostgresClient: and not self.unlimited_query: postgreSQL_pool.putconn(self.connection) + def __execute(self, query, vars=None): + try: + result = self.cursor.cursor_execute(query=query, vars=vars) + except psycopg2.Error as error: + logging.error(f"!!! Error of type:{type(error)} while executing query:") + logging.error(query) + logging.info("starting rollback to allow future execution") + self.connection.rollback() + raise error + return result + def recreate_cursor(self, rollback=False): if rollback: try: