Extract ExecuteContext as in/out argument

So that we can pass out the parsed capabilities to control
retries. This also allows further code optimization.

edgedb/_testbase.py

@@ -372,10 +372,16 @@ def make_test_client(
         database='edgedb',
         user='edgedb',
         password='test',
+        host=...,
+        port=...,
         connection_class=...,
     ):
         conargs = cls.get_connect_args(
             cluster=cluster, database=database, user=user, password=password)
+        if host is not ...:
+            conargs['host'] = host
+        if port is not ...:
+            conargs['port'] = port
         if connection_class is ...:
             connection_class = (
                 asyncio_client.AsyncIOConnection

edgedb/abstract.py

@@ -65,12 +65,42 @@ class QueryContext(typing.NamedTuple):
     retry_options: typing.Optional[options.RetryOptions]
     state: typing.Optional[options.State]
 
+    def lower(
+        self, *, allow_capabilities: enums.Capability
+    ) -> protocol.ExecuteContext:
+        return protocol.ExecuteContext(
+            query=self.query.query,
+            args=self.query.args,
+            kwargs=self.query.kwargs,
+            reg=self.cache.codecs_registry,
+            qc=self.cache.query_cache,
+            output_format=self.query_options.output_format,
+            expect_one=self.query_options.expect_one,
+            required_one=self.query_options.required_one,
+            allow_capabilities=allow_capabilities,
+            state=self.state.as_dict() if self.state else None,
+        )
+
 
 class ExecuteContext(typing.NamedTuple):
     query: QueryWithArgs
     cache: QueryCache
     state: typing.Optional[options.State]
 
+    def lower(
+        self, *, allow_capabilities: enums.Capability
+    ) -> protocol.ExecuteContext:
+        return protocol.ExecuteContext(
+            query=self.query.query,
+            args=self.query.args,
+            kwargs=self.query.kwargs,
+            reg=self.cache.codecs_registry,
+            qc=self.cache.query_cache,
+            output_format=protocol.OutputFormat.NONE,
+            allow_capabilities=allow_capabilities,
+            state=self.state.as_dict() if self.state else None,
+        )
+
 
 @dataclasses.dataclass
 class DescribeContext:

edgedb/base_client.py

@@ -183,17 +183,7 @@ async def privileged_execute(
             )
         else:
             await self._protocol.execute(
-                query=execute_context.query.query,
-                args=execute_context.query.args,
-                kwargs=execute_context.query.kwargs,
-                reg=execute_context.cache.codecs_registry,
-                qc=execute_context.cache.query_cache,
-                output_format=protocol.OutputFormat.NONE,
-                allow_capabilities=enums.Capability.ALL,
-                state=(
-                    execute_context.state.as_dict()
-                    if execute_context.state else None
-                ),
+                execute_context.lower(allow_capabilities=enums.Capability.ALL)
             )
 
     def is_in_transaction(self) -> bool:
@@ -211,56 +201,31 @@ async def raw_query(self, query_context: abstract.QueryContext):
             await self.connect()
 
         reconnect = False
-        capabilities = None
         i = 0
-        args = dict(
-            query=query_context.query.query,
-            args=query_context.query.args,
-            kwargs=query_context.query.kwargs,
-            reg=query_context.cache.codecs_registry,
-            qc=query_context.cache.query_cache,
-            output_format=query_context.query_options.output_format,
-            expect_one=query_context.query_options.expect_one,
-            required_one=query_context.query_options.required_one,
-        )
         if self._protocol.is_legacy:
-            args["allow_capabilities"] = enums.Capability.LEGACY_EXECUTE
+            allow_capabilities = enums.Capability.LEGACY_EXECUTE
         else:
-            args["allow_capabilities"] = enums.Capability.EXECUTE
-            if query_context.state is not None:
-                args["state"] = query_context.state.as_dict()
+            allow_capabilities = enums.Capability.EXECUTE
+        ctx = query_context.lower(allow_capabilities=allow_capabilities)
         while True:
             i += 1
             try:
                 if reconnect:
                     await self.connect(single_attempt=True)
                 if self._protocol.is_legacy:
-                    return await self._protocol.legacy_execute_anonymous(
-                        **args
-                    )
+                    return await self._protocol.legacy_execute_anonymous(ctx)
                 else:
-                    return await self._protocol.query(**args)
+                    return await self._protocol.query(ctx)
             except errors.EdgeDBError as e:
                 if query_context.retry_options is None:
                     raise
                 if not e.has_tag(errors.SHOULD_RETRY):
                     raise e
-                if capabilities is None:
-                    cache_item = query_context.cache.query_cache.get(
-                        query_context.query.query,
-                        query_context.query_options.output_format,
-                        implicit_limit=0,
-                        inline_typenames=False,
-                        inline_typeids=False,
-                        expect_one=query_context.query_options.expect_one,
-                    )
-                    if cache_item is not None:
-                        _, _, _, capabilities = cache_item
                 # A query is read-only if it has no capabilities i.e.
                 # capabilities == 0. Read-only queries are safe to retry.
                 # Explicit transaction conflicts as well.
                 if (
-                    capabilities != 0
+                    ctx.capabilities != 0
                     and not isinstance(e, errors.TransactionConflictError)
                 ):
                     raise e
@@ -281,17 +246,9 @@ async def _execute(self, execute_context: abstract.ExecuteContext) -> None:
             )
         else:
             await self._protocol.execute(
-                query=execute_context.query.query,
-                args=execute_context.query.args,
-                kwargs=execute_context.query.kwargs,
-                reg=execute_context.cache.codecs_registry,
-                qc=execute_context.cache.query_cache,
-                output_format=protocol.OutputFormat.NONE,
-                allow_capabilities=enums.Capability.EXECUTE,
-                state=(
-                    execute_context.state.as_dict()
-                    if execute_context.state else None
-                ),
+                execute_context.lower(
+                    allow_capabilities=enums.Capability.EXECUTE
+                )
             )
 
     async def describe(

edgedb/protocol/protocol.pxd

@@ -78,6 +78,30 @@ cdef class QueryCodecsCache:
              BaseCodec in_type, BaseCodec out_type, int capabilities)
 
 
+cdef class ExecuteContext:
+    cdef:
+        # Input arguments
+        str query
+        object args
+        object kwargs
+        CodecsRegistry reg
+        QueryCodecsCache qc
+        OutputFormat output_format
+        bint expect_one
+        bint required_one
+        int implicit_limit
+        bint inline_typenames
+        bint inline_typeids
+        uint64_t allow_capabilities
+        object state
+
+        # Contextual variables
+        bytes cardinality
+        BaseCodec in_dc
+        BaseCodec out_dc
+        readonly uint64_t capabilities
+
+
 cdef class SansIOProtocol:
 
     cdef: