Postgres pool close finally

credenza/api/session/storage/backends/postgresql.py

@@ -91,11 +91,11 @@ def _get_conn(self):
         conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
         return conn
 
-    def _put_conn(self, conn):
+    def _put_conn(self, conn, close=False):
         if conn is not None:
             if self.trace:
-                logger.debug(f"Returning connection to pool dsn={conn.dsn} status={conn.status}")
-            self.pool.putconn(conn)
+                logger.debug(f"Returning connection to pool dsn={conn.dsn} status={conn.status} close={close}")
+            self.pool.putconn(conn, close=close)
 
     def close(self):
         """
@@ -111,13 +111,24 @@ def close(self):
     def _pooled_execute_stmt(self, sql, params, resultfunc=lambda cur: None):
         """Execute and commit one statement on a pooled connection, returning result of resultfunc applied to cursor.
         """
-        conn = self._get_conn()
-        with conn.cursor() as cur:
-            cur.execute(sql, params)
-            result = resultfunc(cur)
-            conn.commit()
-        self._put_conn(conn)
-        return result
+        conn = None
+        error = 'unknown error'
+        try:
+            conn = self._get_conn()
+            with conn.cursor() as cur:
+                cur.execute(sql, params)
+                result = resultfunc(cur)
+                conn.commit()
+            self._put_conn(conn)
+            conn = None
+            return result
+        except Exception as e:
+            error = e
+            raise
+        finally:
+            if conn is not None:
+                logger.warning(f"Closing pooled connection due to error={error}")
+                self._put_conn(conn, close=True)
 
     def setex(self, key: str, value: Union[str, bytes], ttl: int) -> None:
         expires_at = time.time() + ttl