Fix codegen due to internal API change (#509)

Author: fantix

.github/workflows/tests.yml

@@ -26,7 +26,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
-        edgedb-version: [stable , nightly]
+        edgedb-version: [5, nightly]
         os: [ubuntu-latest, macos-latest, windows-2019]
         loop: [asyncio, uvloop]
         exclude:
@@ -88,6 +88,11 @@ jobs:
         LOOP_IMPL: ${{ matrix.loop }}
         SERVER_VERSION: ${{ matrix.edgedb-version }}
       run: |
+        if [ "${SERVER_VERSION}" = "nightly" ]; then
+          export EDGEDB_TEST_CODEGEN_ASSERT_SUFFIX=.assert
+        else
+          export EDGEDB_TEST_CODEGEN_ASSERT_SUFFIX=.assert${SERVER_VERSION}
+        fi
         if [ "${LOOP_IMPL}" = "uvloop" ]; then
             env USE_UVLOOP=1 python -m unittest -v tests.suite
         else

edgedb/abstract.py

@@ -110,6 +110,22 @@ class DescribeContext:
     output_format: protocol.OutputFormat
     expect_one: bool
 
+    def lower(
+        self, *, allow_capabilities: enums.Capability
+    ) -> protocol.ExecuteContext:
+        return protocol.ExecuteContext(
+            query=self.query,
+            args=None,
+            kwargs=None,
+            reg=protocol.CodecsRegistry(),
+            qc=protocol.LRUMapping(maxsize=1),
+            output_format=self.output_format,
+            expect_one=self.expect_one,
+            inline_typenames=self.inject_type_names,
+            allow_capabilities=allow_capabilities,
+            state=self.state.as_dict() if self.state else None,
+        )
+
 
 @dataclasses.dataclass
 class DescribeResult:

edgedb/base_client.py

@@ -255,23 +255,15 @@ async def _execute(self, execute_context: abstract.ExecuteContext) -> None:
     async def describe(
         self, describe_context: abstract.DescribeContext
     ) -> abstract.DescribeResult:
-        cardinality, in_dc, out_dc, capabilities = await self._protocol._parse(
-            describe_context.query,
-            reg=protocol.CodecsRegistry(),
-            inline_typenames=describe_context.inject_type_names,
-            output_format=describe_context.output_format,
-            expect_one=describe_context.expect_one,
-            allow_capabilities=enums.Capability.EXECUTE,
-            state=(
-                describe_context.state.as_dict()
-                if describe_context.state else None
-            ),
+        ctx = describe_context.lower(
+            allow_capabilities=enums.Capability.EXECUTE
         )
+        await self._protocol._parse(ctx)
         return abstract.DescribeResult(
-            input_type=in_dc.make_type(describe_context),
-            output_type=out_dc.make_type(describe_context),
-            output_cardinality=enums.Cardinality(cardinality[0]),
-            capabilities=capabilities,
+            input_type=ctx.in_dc.make_type(describe_context),
+            output_type=ctx.out_dc.make_type(describe_context),
+            output_cardinality=enums.Cardinality(ctx.cardinality[0]),
+            capabilities=ctx.capabilities,
         )
 
     def terminate(self):

edgedb/protocol/protocol.pxd

@@ -85,9 +85,9 @@ cdef class ExecuteContext:
         object state
 
         # Contextual variables
-        bytes cardinality
-        BaseCodec in_dc
-        BaseCodec out_dc
+        readonly bytes cardinality
+        readonly BaseCodec in_dc
+        readonly BaseCodec out_dc
         readonly uint64_t capabilities
 
     cdef inline bint has_na_cardinality(self)

tests/codegen/test-project1/generated_async_edgeql.py.assert3

@@ -0,0 +1 @@
+generated_async_edgeql.py.assert5

tests/codegen/test-project1/generated_async_edgeql.py.assert5

@@ -0,0 +1,59 @@
+# AUTOGENERATED FROM:
+#     'select_optional_json.edgeql'
+#     'select_scalar.edgeql'
+#     'linked/test_linked.edgeql'
+# WITH:
+#     $ edgedb-py --target async --file --no-skip-pydantic-validation
+
+
+from __future__ import annotations
+import dataclasses
+import edgedb
+import uuid
+
+
+@dataclasses.dataclass
+class SelectOptionalJsonResultItem:
+    id: uuid.UUID
+    snake_case: SelectOptionalJsonResultItemSnakeCase | None
+
+
+@dataclasses.dataclass
+class SelectOptionalJsonResultItemSnakeCase:
+    id: uuid.UUID
+
+
+async def select_optional_json(
+    executor: edgedb.AsyncIOExecutor,
+    arg0: str | None,
+) -> list[tuple[str, SelectOptionalJsonResultItem]]:
+    return await executor.query(
+        """\
+        create type TestCase {
+            create link snake_case -> TestCase;
+        };
+
+        select (<optional json>$0, TestCase {snake_case});\
+        """,
+        arg0,
+    )
+
+
+async def select_scalar(
+    executor: edgedb.AsyncIOExecutor,
+) -> int:
+    return await executor.query_single(
+        """\
+        select 1;\
+        """,
+    )
+
+
+async def test_linked(
+    executor: edgedb.AsyncIOExecutor,
+) -> int:
+    return await executor.query_single(
+        """\
+        select 42\
+        """,
+    )

tests/codegen/test-project2/generated_async_edgeql.py.assert

@@ -22,8 +22,8 @@ import typing
 import uuid
 
 
-Input = typing.Sequence[float]
 MyScalar = int
+V3 = typing.Sequence[float]
 
 
 @dataclasses.dataclass
@@ -61,7 +61,6 @@ class MyEnum(enum.Enum):
 
 @dataclasses.dataclass
 class MyQueryResult:
-    id: uuid.UUID
     a: uuid.UUID
     b: uuid.UUID | None
     c: str
@@ -122,7 +121,6 @@ class MyQueryResult:
 
 @dataclasses.dataclass
 class SelectArgsResult:
-    id: uuid.UUID
     Str: str
     DateTime: datetime.datetime
 
@@ -145,7 +143,7 @@ class SelectObjectResultParamsItem:
 async def custom_vector_input(
     executor: edgedb.AsyncIOExecutor,
     *,
-    input: Input | None = None,
+    input: V3 | None = None,
 ) -> int | None:
     return await executor.query_single(
         """\