Use CompilationRequest for compile_in_tx also (#6896)

edb/server/compiler/compiler.py

@@ -925,6 +925,37 @@ def compile(
         else:
             return unit_group, None
 
+    def compile_in_tx_request(
+        self,
+        state: dbstate.CompilerConnectionState,
+        txid: int,
+        serialized_request: bytes,
+        original_query: str,
+        expect_rollback: bool = False,
+    ) -> Tuple[
+        dbstate.QueryUnitGroup, Optional[dbstate.CompilerConnectionState]
+    ]:
+        request = rpc.CompilationRequest(
+            self.state.compilation_config_serializer
+        )
+        request.deserialize(serialized_request, original_query)
+
+        units, cstate = self.compile_in_tx(
+            state,
+            txid,
+            request.source,
+            request.output_format,
+            request.expect_one,
+            request.implicit_limit,
+            request.inline_typeids,
+            request.inline_typenames,
+            request.protocol_version,
+            request.inline_objectids,
+            request.json_parameters,
+            expect_rollback=expect_rollback,
+        )
+        return units, cstate
+
     def compile_in_tx(
         self,
         state: dbstate.CompilerConnectionState,

edb/server/compiler_pool/worker.py

@@ -200,7 +200,7 @@ def compile_in_tx(cstate, *args, **kwargs):
         cstate = LAST_STATE
     else:
         cstate = pickle.loads(cstate)
-    units, cstate = COMPILER.compile_in_tx(cstate, *args, **kwargs)
+    units, cstate = COMPILER.compile_in_tx_request(cstate, *args, **kwargs)
     LAST_STATE = cstate
     return units, pickle.dumps(cstate, -1)
 

edb/server/dbview/dbview.pyx

@@ -1018,15 +1018,8 @@ cdef class DatabaseConnectionView:
                     self.txid,
                     self._last_comp_state,
                     self._last_comp_state_id,
-                    query_req.source,
-                    query_req.output_format,
-                    query_req.expect_one,
-                    query_req.implicit_limit,
-                    query_req.inline_typeids,
-                    query_req.inline_typenames,
-                    self._protocol_version,
-                    query_req.inline_objectids,
-                    query_req.json_parameters,
+                    query_req.serialize(),
+                    query_req.source.text(),
                     self.in_tx_error(),
                 )
             else:

tests/test_server_compiler.py

@@ -34,6 +34,7 @@
 from edb.testbase import server as tbs
 from edb.server import args as edbargs
 from edb.server import compiler as edbcompiler
+from edb.server.compiler import rpc
 from edb.server import config
 from edb.server.compiler_pool import amsg
 from edb.server.compiler_pool import pool
@@ -422,18 +423,33 @@ async def _test_pool_disconnect_queue(self, pool_class):
                 os.kill(w2.get_pid(), signal.SIGTERM)
                 await asyncio.wait_for(pool_._ready_evt.wait(), LONG_WAIT)
 
+                compiler = edbcompiler.new_compiler(
+                    std_schema=self._std_schema,
+                    reflection_schema=self._refl_schema,
+                    schema_class_layout=self._schema_class_layout,
+                )
+
                 context = edbcompiler.new_compiler_context(
-                    compiler_state=None,
+                    compiler_state=compiler.state,
                     user_schema=self._std_schema,
                     modaliases={None: 'default'},
                 )
+
+                orig_query = 'SELECT 123'
+                request = rpc.CompilationRequest(
+                    compiler.state.compilation_config_serializer
+                ).update(
+                    source=edgeql.Source.from_string(orig_query),
+                    protocol_version=(1, 0),
+                    implicit_limit=101,
+                )
+
                 await asyncio.gather(*(pool_.compile_in_tx(
                     context.state.current_tx().id,
                     pickle.dumps(context.state),
                     0,
-                    edgeql.Source.from_string('SELECT 123'),
-                    edbcompiler.OutputFormat.BINARY,
-                    False, 101, False, True, (1, 0), True
+                    request.serialize(),
+                    orig_query,
                 ) for _ in range(4)))
             finally:
                 await pool_.stop()