Update helper/subscriptions.py

helper/subscriptions.py

@@ -129,37 +129,46 @@ async def get_conn():
     return conn
 
 
+import psycopg
+from psycopg.rows import dict_row
+
 async def execute_query(query: str, params=(), *, fetchone: bool = False, commit: bool = False):
     async with conn_lock:
         connection = await get_conn()
 
         try:
-            with connection.cursor() as cur:
-                cur.execute(query, params)
-                if cur.description is None:
-                    result = None if fetchone else []
-                else:
-                    result = cur.fetchone() if fetchone else cur.fetchall()
-                if commit:
-                    connection.commit()
-                return result
-        except psycopg.OperationalError:
-            connection = psycopg.connect(
+            # Use connection.transaction() to safely manage rollback/commit automatically
+            async with connection.transaction():
+                async with connection.cursor() as cur:
+                    await cur.execute(query, params)
+                    
+                    if cur.description is None:
+                        return None if fetchone else []
+                        
+                    return await cur.fetchone() if fetchone else await cur.fetchall()
+
+        except (psycopg.OperationalError, psycopg.InterfaceError):
+            # 1. Handle dropped connections by reconnecting asynchronously
+            connection = await psycopg.AsyncConnection.connect(
                 POSTGRE_SECRET,
                 row_factory=dict_row,
                 sslmode="verify-full",
                 sslrootcert="prod-ca-2021.crt",
             )
             globals()["conn"] = connection
-            with connection.cursor() as cur:
-                cur.execute(query, params)
-                if cur.description is None:
-                    result = None if fetchone else []
-                else:
-                    result = cur.fetchone() if fetchone else cur.fetchall()
-                if commit:
-                    connection.commit()
-                return result
+            
+            # 2. Retry the query safely on the new connection
+            async with connection.transaction():
+                async with connection.cursor() as cur:
+                    await cur.execute(query, params)
+                    if cur.description is None:
+                        return None if fetchone else []
+                    return await cur.fetchone() if fetchone else await cur.fetchall()
+                    
+        except psycopg.Error:
+            # 3. Handle bad SQL/transactions: the transaction context manager 
+            # automatically issued a ROLLBACK before reaching this point.
+            raise
 
 
 def normalize_plan_key(plan_name: str | None) -> str: