Repository: geldata/gel
Pull request: Supress idle tx timeout during compilation
I tested this by adding a sleep when the variable `slowcompile` is used. 
Fixes #6758.

Problem statement:
We should suspend the pg transaction timeout while waiting for the compiler
If there is a long running compile job (probably computing a migration?), we can spuriously trigger an idle-in-transaction timeout, since even though the transaction is not idle from edgedb's perspective, it *is* from postgres's.

Hints from issue discussion:
Another option would be to have a heartbeat task that sends keep alive queries. That might actually be better, since in the common case it would never have to do anything.

Commit message: Suppress the timeout instead

Patch to review:
diff --git a/edb/server/dbview/dbview.pyx b/edb/server/dbview/dbview.pyx
index d949d8d64c6..296304ec0fd 100644
--- a/edb/server/dbview/dbview.pyx
+++ b/edb/server/dbview/dbview.pyx
@@ -977,6 +977,9 @@ cdef class DatabaseConnectionView:
         query_req: QueryRequestInfo,
         cached_globally=False,
         use_metrics=True,
+        # allow passing in the query_unit_group, in case the call site needs
+        # to make decisions based on whether the cache is hit
+        query_unit_group=None,
     ) -> CompiledQuery:
         source = query_req.source
         if cached_globally:
@@ -990,7 +993,7 @@ cdef class DatabaseConnectionView:
                 if self._query_cache_enabled
                 else None
             )
-        else:
+        elif query_unit_group is None:
             query_unit_group = self.lookup_compiled_query(query_req)
         cached = True
         if query_unit_group is None:
diff --git a/edb/server/protocol/binary.pyx b/edb/server/protocol/binary.pyx
index 04dca738cc3..b415806ca53 100644
--- a/edb/server/protocol/binary.pyx
+++ b/edb/server/protocol/binary.pyx
@@ -499,6 +499,27 @@ cdef class EdgeConnection(frontend.FrontendConnection):
         else:
             return edgeql.NormalizedSource.from_string(text)
 
+    async def _suppress_tx_timeout(self):
+        async with self.with_pgcon() as conn:
+            await conn.sql_execute(b'''
+                select pg_catalog.set_config(
+                    'idle_in_transaction_session_timeout', '0', false)
+            ''')
+
+    async def _restore_tx_timeout(self, dbview.DatabaseConnectionView dbv):
+        old_timeout = dbv.get_session_config().get(
+            'session_idle_transaction_timeout',
+        )
+        timeout = (
+            'NULL' if not old_timeout
+            else repr(old_timeout.value.to_backend_str())
+        )
+        async with self.with_pgcon() as conn:
+            await conn.sql_execute(f'''
+                select pg_catalog.set_config(
+                    'idle_in_transaction_session_timeout', {timeout}, false)
+            '''.encode('utf-8'))
+
     async def _parse(
         self,
         dbview.QueryRequestInfo query_req,
@@ -526,7 +547,21 @@ cdef class EdgeConnection(frontend.FrontendConnection):
             self.debug_print('Extra variables', source.variables(),
                              'after', source.first_extra())
 
-        return await dbv.parse(query_req)
+        query_unit_group = dbv.lookup_compiled_query(query_req)
+
+        # If we have to do a compile within a transaction, suppress
+        # the idle_in_transaction_session_timeout.
+        suppress_timeout = (
+            dbv.in_tx() and not dbv.in_tx_error() and query_unit_group is None
+        )
+        if suppress_timeout:
+            await self._suppress_tx_timeout()
+
+        try:
+            return await dbv.parse(query_req, query_unit_group=query_unit_group)
+        finally:
+            if suppress_timeout:
+                await self._restore_tx_timeout(dbv)
 
     cdef parse_cardinality(self, bytes card):
         if card[0] == CARD_MANY.value:
diff --git a/edb/server/protocol/frontend.pyx b/edb/server/protocol/frontend.pyx
index f812a328389..c93c7c34605 100644
--- a/edb/server/protocol/frontend.pyx
+++ b/edb/server/protocol/frontend.pyx
@@ -18,6 +18,7 @@
 
 
 import asyncio
+import contextlib
 import logging
 import time
 
@@ -167,6 +168,14 @@ cdef class FrontendConnection(AbstractFrontendConnection):
                     discard=debug.flags.server_clobber_pg_conns,
                 )
 
+    @contextlib.asynccontextmanager
+    async def with_pgcon(self):
+        con = await self.get_pgcon()
+        try:
+            yield con
+        finally:
+            self.maybe_release_pgcon(con)
+
     def on_aborted_pgcon(self, pgcon.PGConnection conn):
         try:
             self._pinned_pgcon = None


Code review comments:
diff_hunk: @@ -499,6 +499,27 @@ cdef class EdgeConnection(frontend.FrontendConnection):
         else:
             return edgeql.NormalizedSource.from_string(text)
 
+    async def _suppress_tx_timeout(self):
+        async with self.with_pgcon() as conn:
+            await conn.sql_execute(b'''
+                select pg_catalog.set_config(
+                    'idle_in_transaction_session_timeout', '0', false)
path: edb/server/protocol/binary.pyx
text: @user1:
You probably want the opposite?

```suggestion
                    'idle_in_transaction_session_timeout', '0', true)
```

@author:
Sure. I don't think it should matter in practice, since there shouldn't be any way for this set_config to ever get committed