From 8beae3188915a525e7fc2a23971deb42517961b1 Mon Sep 17 00:00:00 2001
From: Taha Yassine Kraiem <tahayk2@gmail.com>
Date: Thu, 9 Mar 2023 18:02:46 +0100
Subject: [PATCH] feat(chalice): pg execute wrapper to handle all query
 failures

---
 api/chalicelib/utils/pg_client.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

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: