Repository: edgedb/edgedb-python (language: python)

Issue:
Extract ExecuteContext as in/out argument
So that we can pass out the parsed capabilities to control retries. This also allows further code optimization.

Fixes #493 
protocol _parse broken
Hello,

it seems that edgedb-python is broken is broken due to arguments changed in last version : 


```
[crunch@wardesk musicbot (master)]$ poetry run edgedb-py -I musicbot-prod --dir musicbot/queries
Found EdgeDB project: /home/crunch/projects/musicbot
Processing /home/crunch/projects/musicbot/musicbot/queries/select_folder.edgeql
Traceback (most recent call last):
  File "/home/crunch/.cache/pypoetry/virtualenvs/musicbot-Fs1OZ5tc-py3.11/bin/edgedb-py", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/home/crunch/.cache/pypoetry/virtualenvs/musicbot-Fs1OZ5tc-py3.11/lib/python3.11/site-packages/edgedb/codegen/cli.py", line 99, in main
    generator.Generator(args).run()
  File "/home/crunch/.cache/pypoetry/virtualenvs/musicbot-Fs1OZ5tc-py3.11/lib/python3.11/site-packages/edgedb/codegen/generator.py", line 212, in run
    self._process_dir(search_dir)
  File "/home/crunch/.cache/pypoetry/virtualenvs/musicbot-Fs1OZ5tc-py3.11/lib/python3.11/site-packages/edgedb/codegen/generator.py", line 235, in _process_dir
    self._process_file(file_or_dir)
  File "/home/crunch/.cache/pypoetry/virtualenvs/musicbot-Fs1OZ5tc-py3.11/lib/python3.11/site-packages/edgedb/codegen/generator.py", line 247, in _process_file
    dr = self._client._describe_query(query, inject_type_names=True)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/crunch/.cache/pypoetry/virtualenvs/musicbot-Fs1OZ5tc-py3.11/lib/python3.11/site-packages/edgedb/blocking_client.py", line 441, in _describe_query
    return self._iter_coroutine(self._describe(abstract.DescribeContext(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/crunch/.cache/pypoetry/virtualenvs/musicbot-Fs1OZ5tc-py3.11/lib/python3.11/site-packages/edgedb/blocking_client.py", line 394, in _iter_coroutine
    coro.send(None)
  File "/home/crunch/.cache/pypoetry/virtualenvs/musicbot-Fs1OZ5tc-py3.11/lib/python3.11/site-packages/edgedb/base_client.py", line 726, in _describe
    return await con.describe(describe_context)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/crunch/.cache/pypoetry/virtualenvs/musicbot-Fs1OZ5tc-py3.11/lib/python3.11/site-packages/edgedb/base_client.py", line 258, in describe
    cardinality, in_dc, out_dc, capabilities = await self._protocol._parse(
                                                     ^^^^^^^^^^^^^^^^^^^^^^
TypeError: SansIOProtocol._parse() takes no keyword arguments
```


Discussion / hints:

I experience this too in version 2.0.0. 

```
Found EdgeDB project: /Users/default/work/code/learn/edgedb/quickstart
Processing /Users/default/work/code/learn/edgedb/quickstart/apps/python_client/queries/get_user_by_name.edgeql
Traceback (most recent call last):
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/bin/edgedb-py", line 8, in <module>
    sys.exit(main())
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/codegen/cli.py", line 99, in main
    generator.Generator(args).run()
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/codegen/generator.py", line 214, in run
    self._process_dir(self._project_dir)
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/codegen/generator.py", line 233, in _process_dir
    self._process_dir(file_or_dir)
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/codegen/generator.py", line 233, in _process_dir
    self._process_dir(file_or_dir)
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/codegen/generator.py", line 233, in _process_dir
    self._process_dir(file_or_dir)
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/codegen/generator.py", line 235, in _process_dir
    self._process_file(file_or_dir)
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/codegen/generator.py", line 247, in _process_file
    dr = self._client._describe_query(query, inject_type_names=True)
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/blocking_client.py", line 441, in _describe_query
    return self._iter_coroutine(self._describe(abstract.DescribeContext(
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/blocking_client.py", line 394, in _iter_coroutine
    coro.send(None)
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/base_client.py", line 726, in _describe
    return await con.describe(describe_context)
  File "/Users/default/work/code/learn/edgedb/quickstart/apps/python_client/env/lib/python3.10/site-packages/edgedb/base_client.py", line 258, in describe
    cardinality, in_dc, out_dc, capabilities = await self._protocol._parse(
TypeError: BlockingIOProtocol._parse() takes no keyword arguments
```
Downgrading to 1.9.0 fixes it.



Fix patch:
diff --git a/edgedb/abstract.py b/edgedb/abstract.py
index 860275c7..dbac4503 100644
--- a/edgedb/abstract.py
+++ b/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:
diff --git a/edgedb/base_client.py b/edgedb/base_client.py
index 0744211c..382da052 100644
--- a/edgedb/base_client.py
+++ b/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):
diff --git a/edgedb/protocol/protocol.pxd b/edgedb/protocol/protocol.pxd
index 402c988c..7140ae4e 100644
--- a/edgedb/protocol/protocol.pxd
+++ b/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)


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 28cb0d53..e91e718a 100644
--- a/.github/workflows/tests.yml
+++ b/.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
diff --git a/tests/codegen/test-project1/generated_async_edgeql.py.assert3 b/tests/codegen/test-project1/generated_async_edgeql.py.assert3
new file mode 120000
index 00000000..cb37cab3
--- /dev/null
+++ b/tests/codegen/test-project1/generated_async_edgeql.py.assert3
@@ -0,0 +1,1 @@
+generated_async_edgeql.py.assert5
\ No newline at end of file
diff --git a/tests/codegen/test-project1/generated_async_edgeql.py.assert5 b/tests/codegen/test-project1/generated_async_edgeql.py.assert5
new file mode 100644
index 00000000..f4e8fea7
--- /dev/null
+++ b/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\
+        """,
+    )
diff --git a/tests/codegen/test-project2/generated_async_edgeql.py.assert b/tests/codegen/test-project2/generated_async_edgeql.py.assert
index 2095cc9a..007f2af1 100644
--- a/tests/codegen/test-project2/generated_async_edgeql.py.assert
+++ b/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(
         """\
diff --git a/tests/codegen/test-project2/generated_async_edgeql.py.assert3 b/tests/codegen/test-project2/generated_async_edgeql.py.assert3
new file mode 100644
index 00000000..2095cc9a
--- /dev/null
+++ b/tests/codegen/test-project2/generated_async_edgeql.py.assert3
@@ -0,0 +1,443 @@
+# AUTOGENERATED FROM:
+#     'scalar/custom_vector_input.edgeql'
+#     'object/link_prop.edgeql'
+#     'parpkg/subpkg/my_query.edgeql'
+#     'argnames/query_one.edgeql'
+#     'parpkg/select_args.edgeql'
+#     'object/select_object.edgeql'
+#     'object/select_objects.edgeql'
+#     'scalar/select_scalar.edgeql'
+#     'scalar/select_scalars.edgeql'
+# WITH:
+#     $ edgedb-py --target async --file --no-skip-pydantic-validation
+
+
+from __future__ import annotations
+import array
+import dataclasses
+import datetime
+import edgedb
+import enum
+import typing
+import uuid
+
+
+Input = typing.Sequence[float]
+MyScalar = int
+
+
+@dataclasses.dataclass
+class LinkPropResult:
+    id: uuid.UUID
+    name: str
+    friends: list[LinkPropResultFriendsItem]
+
+
+@dataclasses.dataclass
+class LinkPropResultFriendsItem:
+    id: uuid.UUID
+    name: str
+    created_at: datetime.datetime | None
+
+    @typing.overload
+    def __getitem__(self, key: typing.Literal["@created_at"]) -> datetime.datetime | None:
+        ...
+
+    @typing.overload
+    def __getitem__(self, key: typing.Literal["@strength"]) -> float | None:
+        ...
+
+    def __getitem__(self, key: str) -> typing.Any:
+        raise NotImplementedError
+
+
+class MyEnum(enum.Enum):
+    THIS = "This"
+    THAT = "That"
+    E_1 = "1"
+    F_B = "f. b"
+    F_B_1 = "f-b"
+
+
+@dataclasses.dataclass
+class MyQueryResult:
+    id: uuid.UUID
+    a: uuid.UUID
+    b: uuid.UUID | None
+    c: str
+    d: str | None
+    e: bytes
+    f: bytes | None
+    g: int
+    h: int | None
+    i: int
+    j: int | None
+    k: int
+    l: int | None
+    m: float
+    n: float | None
+    o: float
+    p: float | None
+    q: bool
+    r: bool | None
+    s: datetime.datetime
+    t: datetime.datetime | None
+    u: datetime.datetime
+    v: datetime.datetime | None
+    w: datetime.date
+    x: datetime.date | None
+    y: datetime.time
+    z: datetime.time | None
+    aa: datetime.timedelta
+    ab: datetime.timedelta | None
+    ac: int
+    ad: int | None
+    ae: edgedb.RelativeDuration
+    af: edgedb.RelativeDuration | None
+    ag: edgedb.DateDuration
+    ah: edgedb.DateDuration | None
+    ai: edgedb.ConfigMemory
+    aj: edgedb.ConfigMemory | None
+    ak: edgedb.Range[int]
+    al: edgedb.Range[int] | None
+    am: edgedb.Range[int]
+    an: edgedb.Range[int] | None
+    ao: edgedb.Range[float]
+    ap: edgedb.Range[float] | None
+    aq: edgedb.Range[float]
+    ar: edgedb.Range[float] | None
+    as_: edgedb.Range[datetime.datetime]
+    at: edgedb.Range[datetime.datetime] | None
+    au: edgedb.Range[datetime.datetime]
+    av: edgedb.Range[datetime.datetime] | None
+    aw: edgedb.Range[datetime.date]
+    ax: edgedb.Range[datetime.date] | None
+    ay: MyScalar
+    az: MyScalar | None
+    ba: MyEnum
+    bb: MyEnum | None
+    bc: array.array
+    bd: array.array | None
+
+
+@dataclasses.dataclass
+class SelectArgsResult:
+    id: uuid.UUID
+    Str: str
+    DateTime: datetime.datetime
+
+
+@dataclasses.dataclass
+class SelectObjectResult:
+    id: uuid.UUID
+    Name: str
+    Language: str
+    Params: list[SelectObjectResultParamsItem]
+
+
+@dataclasses.dataclass
+class SelectObjectResultParamsItem:
+    id: uuid.UUID
+    Name: str
+    Default: str | None
+
+
+async def custom_vector_input(
+    executor: edgedb.AsyncIOExecutor,
+    *,
+    input: Input | None = None,
+) -> int | None:
+    return await executor.query_single(
+        """\
+        select 42 filter exists <optional v3>$input;\
+        """,
+        input=input,
+    )
+
+
+async def link_prop(
+    executor: edgedb.AsyncIOExecutor,
+) -> list[LinkPropResult]:
+    return await executor.query(
+        """\
+        create type Person {
+            create required property name -> str;
+            create property created_at -> datetime;
+            create multi link friends -> Person {
+                create property strength -> float64;
+                create property created_at -> datetime;
+            }
+        };
+
+        select Person {
+            name,
+            friends: {
+                name,
+                @created_at,
+                created_at,
+                @strength,
+            }
+        };\
+        """,
+    )
+
+
+async def my_query(
+    executor: edgedb.AsyncIOExecutor,
+    *,
+    a: uuid.UUID,
+    b: uuid.UUID | None = None,
+    c: str,
+    d: str | None = None,
+    e: bytes,
+    f: bytes | None = None,
+    g: int,
+    h: int | None = None,
+    i: int,
+    j: int | None = None,
+    k: int,
+    l: int | None = None,
+    m: float,
+    n: float | None = None,
+    o: float,
+    p: float | None = None,
+    q: bool,
+    r: bool | None = None,
+    s: datetime.datetime,
+    t: datetime.datetime | None = None,
+    u: datetime.datetime,
+    v: datetime.datetime | None = None,
+    w: datetime.date,
+    x: datetime.date | None = None,
+    y: datetime.time,
+    z: datetime.time | None = None,
+    aa: datetime.timedelta,
+    ab: datetime.timedelta | None = None,
+    ac: int,
+    ad: int | None = None,
+    ae: edgedb.RelativeDuration,
+    af: edgedb.RelativeDuration | None = None,
+    ag: edgedb.DateDuration,
+    ah: edgedb.DateDuration | None = None,
+    ai: edgedb.ConfigMemory,
+    aj: edgedb.ConfigMemory | None = None,
+    ak: edgedb.Range[int],
+    al: edgedb.Range[int] | None = None,
+    am: edgedb.Range[int],
+    an: edgedb.Range[int] | None = None,
+    ao: edgedb.Range[float],
+    ap: edgedb.Range[float] | None = None,
+    aq: edgedb.Range[float],
+    ar: edgedb.Range[float] | None = None,
+    as_: edgedb.Range[datetime.datetime],
+    at: edgedb.Range[datetime.datetime] | None = None,
+    au: edgedb.Range[datetime.datetime],
+    av: edgedb.Range[datetime.datetime] | None = None,
+    aw: edgedb.Range[datetime.date],
+    ax: edgedb.Range[datetime.date] | None = None,
+    bc: typing.Sequence[float],
+    bd: typing.Sequence[float] | None = None,
+) -> MyQueryResult:
+    return await executor.query_single(
+        """\
+        create scalar type MyScalar extending int64;
+        create scalar type MyEnum extending enum<'This', 'That', '1', 'f. b', 'f-b'>;
+
+        select {
+        	a := <uuid>$a,
+        	b := <optional uuid>$b,
+        	c := <str>$c,
+        	d := <optional str>$d,
+        	e := <bytes>$e,
+        	f := <optional bytes>$f,
+        	g := <int16>$g,
+        	h := <optional int16>$h,
+        	i := <int32>$i,
+        	j := <optional int32>$j,
+        	k := <int64>$k,
+        	l := <optional int64>$l,
+        	m := <float32>$m,
+        	n := <optional float32>$n,
+        	o := <float64>$o,
+        	p := <optional float64>$p,
+        	q := <bool>$q,
+        	r := <optional bool>$r,
+        	s := <datetime>$s,
+        	t := <optional datetime>$t,
+        	u := <cal::local_datetime>$u,
+        	v := <optional cal::local_datetime>$v,
+        	w := <cal::local_date>$w,
+        	x := <optional cal::local_date>$x,
+        	y := <cal::local_time>$y,
+        	z := <optional cal::local_time>$z,
+        	aa := <duration>$aa,
+        	ab := <optional duration>$ab,
+        	ac := <bigint>$ac,
+        	ad := <optional bigint>$ad,
+        	ae := <cal::relative_duration>$ae,
+        	af := <optional cal::relative_duration>$af,
+        	ag := <cal::date_duration>$ag,
+        	ah := <optional cal::date_duration>$ah,
+        	ai := <cfg::memory>$ai,
+        	aj := <optional cfg::memory>$aj,
+        	ak := <range<int32>>$ak,
+        	al := <optional range<int32>>$al,
+        	am := <range<int64>>$am,
+        	an := <optional range<int64>>$an,
+        	ao := <range<float32>>$ao,
+        	ap := <optional range<float32>>$ap,
+        	aq := <range<float64>>$aq,
+        	ar := <optional range<float64>>$ar,
+        	as_ := <range<datetime>>$as_,
+        	at := <optional range<datetime>>$at,
+        	au := <range<cal::local_datetime>>$au,
+        	av := <optional range<cal::local_datetime>>$av,
+        	aw := <range<cal::local_date>>$aw,
+        	ax := <optional range<cal::local_date>>$ax,
+        	ay := <MyScalar>1,
+        	az := <optional MyScalar>{},
+        	ba := MyEnum.This,
+        	bb := <optional MyEnum>{},
+        	bc := <ext::pgvector::vector>$bc,
+        	bd := <optional ext::pgvector::vector>$bd,
+        }\
+        """,
+        a=a,
+        b=b,
+        c=c,
+        d=d,
+        e=e,
+        f=f,
+        g=g,
+        h=h,
+        i=i,
+        j=j,
+        k=k,
+        l=l,
+        m=m,
+        n=n,
+        o=o,
+        p=p,
+        q=q,
+        r=r,
+        s=s,
+        t=t,
+        u=u,
+        v=v,
+        w=w,
+        x=x,
+        y=y,
+        z=z,
+        aa=aa,
+        ab=ab,
+        ac=ac,
+        ad=ad,
+        ae=ae,
+        af=af,
+        ag=ag,
+        ah=ah,
+        ai=ai,
+        aj=aj,
+        ak=ak,
+        al=al,
+        am=am,
+        an=an,
+        ao=ao,
+        ap=ap,
+        aq=aq,
+        ar=ar,
+        as_=as_,
+        at=at,
+        au=au,
+        av=av,
+        aw=aw,
+        ax=ax,
+        bc=bc,
+        bd=bd,
+    )
+
+
+async def query_one(
+    executor: edgedb.AsyncIOExecutor,
+    *,
+    arg_name_with_underscores: int,
+) -> int:
+    return await executor.query_single(
+        """\
+        select <int64>$arg_name_with_underscores\
+        """,
+        arg_name_with_underscores=arg_name_with_underscores,
+    )
+
+
+async def select_args(
+    executor: edgedb.AsyncIOExecutor,
+    *,
+    arg_str: str,
+    arg_datetime: datetime.datetime,
+) -> SelectArgsResult:
+    return await executor.query_single(
+        """\
+        select {
+        	Str := <str>$arg_str,
+        	DateTime := <datetime>$arg_datetime,
+        }\
+        """,
+        arg_str=arg_str,
+        arg_datetime=arg_datetime,
+    )
+
+
+async def select_object(
+    executor: edgedb.AsyncIOExecutor,
+) -> SelectObjectResult | None:
+    return await executor.query_single(
+        """\
+        select schema::Function {
+          Name := .name,
+          Language := .language,
+          Params := .params {
+            Name := .name,
+            Default := .default,
+          }
+        }
+        limit 1;\
+        """,
+    )
+
+
+async def select_objects(
+    executor: edgedb.AsyncIOExecutor,
+) -> list[SelectObjectResult]:
+    return await executor.query(
+        """\
+        select schema::Function {
+          Name := .name,
+          Language := .language,
+          Params := .params {
+            Name := .name,
+            Default := .default,
+          }
+        }\
+        """,
+    )
+
+
+async def select_scalar(
+    executor: edgedb.AsyncIOExecutor,
+) -> int:
+    return await executor.query_single(
+        """\
+        select 1;\
+        """,
+    )
+
+
+async def select_scalars(
+    executor: edgedb.AsyncIOExecutor,
+) -> list[edgedb.ConfigMemory]:
+    return await executor.query(
+        """\
+        select <cfg::memory>{1, 2, 3};\
+        """,
+    )
diff --git a/tests/codegen/test-project2/generated_async_edgeql.py.assert4 b/tests/codegen/test-project2/generated_async_edgeql.py.assert5
similarity index 100%
rename from tests/codegen/test-project2/generated_async_edgeql.py.assert4
rename to tests/codegen/test-project2/generated_async_edgeql.py.assert5
diff --git a/tests/codegen/test-project2/object/select_object_async_edgeql.py.assert b/tests/codegen/test-project2/object/select_object_async_edgeql.py.assert
index 56d040a2..3e121c7a 100644
--- a/tests/codegen/test-project2/object/select_object_async_edgeql.py.assert
+++ b/tests/codegen/test-project2/object/select_object_async_edgeql.py.assert
@@ -20,7 +20,7 @@ class NoPydanticValidation:
     def __get_validators__(cls):
         # Pydantic 1.x
         from pydantic.dataclasses import dataclass as pydantic_dataclass
-        pydantic_dataclass(cls)
+        _ = pydantic_dataclass(cls)
         cls.__pydantic_model__.__get_validators__ = lambda: []
         return []
 
diff --git a/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert b/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert
index b318bb5c..7643f4f4 100644
--- a/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert
+++ b/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert
@@ -6,7 +6,6 @@ from __future__ import annotations
 import dataclasses
 import datetime
 import edgedb
-import uuid
 
 
 class NoPydanticValidation:
@@ -27,7 +26,6 @@ class NoPydanticValidation:
 
 @dataclasses.dataclass
 class SelectArgsResult(NoPydanticValidation):
-    id: uuid.UUID
     Str: str
     DateTime: datetime.datetime
 
diff --git a/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert3 b/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert3
new file mode 120000
index 00000000..49afd52b
--- /dev/null
+++ b/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert3
@@ -0,0 +1,1 @@
+select_args_async_edgeql.py.assert5
\ No newline at end of file
diff --git a/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert5 b/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert5
new file mode 100644
index 00000000..b318bb5c
--- /dev/null
+++ b/tests/codegen/test-project2/parpkg/select_args_async_edgeql.py.assert5
@@ -0,0 +1,50 @@
+# AUTOGENERATED FROM 'parpkg/select_args.edgeql' WITH:
+#     $ edgedb-py
+
+
+from __future__ import annotations
+import dataclasses
+import datetime
+import edgedb
+import uuid
+
+
+class NoPydanticValidation:
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source_type, _handler):
+        # Pydantic 2.x
+        from pydantic_core.core_schema import any_schema
+        return any_schema()
+
+    @classmethod
+    def __get_validators__(cls):
+        # Pydantic 1.x
+        from pydantic.dataclasses import dataclass as pydantic_dataclass
+        _ = pydantic_dataclass(cls)
+        cls.__pydantic_model__.__get_validators__ = lambda: []
+        return []
+
+
+@dataclasses.dataclass
+class SelectArgsResult(NoPydanticValidation):
+    id: uuid.UUID
+    Str: str
+    DateTime: datetime.datetime
+
+
+async def select_args(
+    executor: edgedb.AsyncIOExecutor,
+    *,
+    arg_str: str,
+    arg_datetime: datetime.datetime,
+) -> SelectArgsResult:
+    return await executor.query_single(
+        """\
+        select {
+        	Str := <str>$arg_str,
+        	DateTime := <datetime>$arg_datetime,
+        }\
+        """,
+        arg_str=arg_str,
+        arg_datetime=arg_datetime,
+    )
diff --git a/tests/codegen/test-project2/parpkg/select_args_edgeql.py.assert b/tests/codegen/test-project2/parpkg/select_args_edgeql.py.assert
index cc0d8cfc..88fd6cae 100644
--- a/tests/codegen/test-project2/parpkg/select_args_edgeql.py.assert
+++ b/tests/codegen/test-project2/parpkg/select_args_edgeql.py.assert
@@ -6,12 +6,10 @@ from __future__ import annotations
 import dataclasses
 import datetime
 import edgedb
-import uuid
 
 
 @dataclasses.dataclass
 class SelectArgsResult:
-    id: uuid.UUID
     Str: str
     DateTime: datetime.datetime
 
diff --git a/tests/codegen/test-project2/parpkg/select_args_edgeql.py.assert3 b/tests/codegen/test-project2/parpkg/select_args_edgeql.py.assert3
new file mode 120000
index 00000000..8a26053a
--- /dev/null
+++ b/tests/codegen/test-project2/parpkg/select_args_edgeql.py.assert3
@@ -0,0 +1,1 @@
+select_args_edgeql.py.assert5
\ No newline at end of file
diff --git a/tests/codegen/test-project2/parpkg/select_args_edgeql.py.assert5 b/tests/codegen/test-project2/parpkg/select_args_edgeql.py.assert5
new file mode 100644
index 00000000..cc0d8cfc
--- /dev/null
+++ b/tests/codegen/test-project2/parpkg/select_args_edgeql.py.assert5
@@ -0,0 +1,34 @@
+# AUTOGENERATED FROM 'parpkg/select_args.edgeql' WITH:
+#     $ edgedb-py --target blocking --no-skip-pydantic-validation
+
+
+from __future__ import annotations
+import dataclasses
+import datetime
+import edgedb
+import uuid
+
+
+@dataclasses.dataclass
+class SelectArgsResult:
+    id: uuid.UUID
+    Str: str
+    DateTime: datetime.datetime
+
+
+def select_args(
+    executor: edgedb.Executor,
+    *,
+    arg_str: str,
+    arg_datetime: datetime.datetime,
+) -> SelectArgsResult:
+    return executor.query_single(
+        """\
+        select {
+        	Str := <str>$arg_str,
+        	DateTime := <datetime>$arg_datetime,
+        }\
+        """,
+        arg_str=arg_str,
+        arg_datetime=arg_datetime,
+    )
diff --git a/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert b/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert
index 05e4a7e2..79b133af 100644
--- a/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert
+++ b/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert
@@ -41,7 +41,6 @@ class MyEnum(enum.Enum):
 
 @dataclasses.dataclass
 class MyQueryResult(NoPydanticValidation):
-    id: uuid.UUID
     a: uuid.UUID
     b: typing.Optional[uuid.UUID]
     c: str
diff --git a/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert3 b/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert3
new file mode 120000
index 00000000..502a7328
--- /dev/null
+++ b/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert3
@@ -0,0 +1,1 @@
+my_query_async_edgeql.py.assert5
\ No newline at end of file
diff --git a/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert5 b/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert5
new file mode 100644
index 00000000..05e4a7e2
--- /dev/null
+++ b/tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert5
@@ -0,0 +1,275 @@
+# AUTOGENERATED FROM 'parpkg/subpkg/my_query.edgeql' WITH:
+#     $ edgedb-py
+
+
+from __future__ import annotations
+import array
+import dataclasses
+import datetime
+import edgedb
+import enum
+import typing
+import uuid
+
+
+MyScalar = int
+
+
+class NoPydanticValidation:
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source_type, _handler):
+        # Pydantic 2.x
+        from pydantic_core.core_schema import any_schema
+        return any_schema()
+
+    @classmethod
+    def __get_validators__(cls):
+        # Pydantic 1.x
+        from pydantic.dataclasses import dataclass as pydantic_dataclass
+        _ = pydantic_dataclass(cls)
+        cls.__pydantic_model__.__get_validators__ = lambda: []
+        return []
+
+
+class MyEnum(enum.Enum):
+    THIS = "This"
+    THAT = "That"
+    E_1 = "1"
+    F_B = "f. b"
+    F_B_1 = "f-b"
+
+
+@dataclasses.dataclass
+class MyQueryResult(NoPydanticValidation):
+    id: uuid.UUID
+    a: uuid.UUID
+    b: typing.Optional[uuid.UUID]
+    c: str
+    d: typing.Optional[str]
+    e: bytes
+    f: typing.Optional[bytes]
+    g: int
+    h: typing.Optional[int]
+    i: int
+    j: typing.Optional[int]
+    k: int
+    l: typing.Optional[int]
+    m: float
+    n: typing.Optional[float]
+    o: float
+    p: typing.Optional[float]
+    q: bool
+    r: typing.Optional[bool]
+    s: datetime.datetime
+    t: typing.Optional[datetime.datetime]
+    u: datetime.datetime
+    v: typing.Optional[datetime.datetime]
+    w: datetime.date
+    x: typing.Optional[datetime.date]
+    y: datetime.time
+    z: typing.Optional[datetime.time]
+    aa: datetime.timedelta
+    ab: typing.Optional[datetime.timedelta]
+    ac: int
+    ad: typing.Optional[int]
+    ae: edgedb.RelativeDuration
+    af: typing.Optional[edgedb.RelativeDuration]
+    ag: edgedb.DateDuration
+    ah: typing.Optional[edgedb.DateDuration]
+    ai: edgedb.ConfigMemory
+    aj: typing.Optional[edgedb.ConfigMemory]
+    ak: edgedb.Range[int]
+    al: typing.Optional[edgedb.Range[int]]
+    am: edgedb.Range[int]
+    an: typing.Optional[edgedb.Range[int]]
+    ao: edgedb.Range[float]
+    ap: typing.Optional[edgedb.Range[float]]
+    aq: edgedb.Range[float]
+    ar: typing.Optional[edgedb.Range[float]]
+    as_: edgedb.Range[datetime.datetime]
+    at: typing.Optional[edgedb.Range[datetime.datetime]]
+    au: edgedb.Range[datetime.datetime]
+    av: typing.Optional[edgedb.Range[datetime.datetime]]
+    aw: edgedb.Range[datetime.date]
+    ax: typing.Optional[edgedb.Range[datetime.date]]
+    ay: MyScalar
+    az: typing.Optional[MyScalar]
+    ba: MyEnum
+    bb: typing.Optional[MyEnum]
+    bc: array.array
+    bd: typing.Optional[array.array]
+
+
+async def my_query(
+    executor: edgedb.AsyncIOExecutor,
+    *,
+    a: uuid.UUID,
+    b: typing.Optional[uuid.UUID] = None,
+    c: str,
+    d: typing.Optional[str] = None,
+    e: bytes,
+    f: typing.Optional[bytes] = None,
+    g: int,
+    h: typing.Optional[int] = None,
+    i: int,
+    j: typing.Optional[int] = None,
+    k: int,
+    l: typing.Optional[int] = None,
+    m: float,
+    n: typing.Optional[float] = None,
+    o: float,
+    p: typing.Optional[float] = None,
+    q: bool,
+    r: typing.Optional[bool] = None,
+    s: datetime.datetime,
+    t: typing.Optional[datetime.datetime] = None,
+    u: datetime.datetime,
+    v: typing.Optional[datetime.datetime] = None,
+    w: datetime.date,
+    x: typing.Optional[datetime.date] = None,
+    y: datetime.time,
+    z: typing.Optional[datetime.time] = None,
+    aa: datetime.timedelta,
+    ab: typing.Optional[datetime.timedelta] = None,
+    ac: int,
+    ad: typing.Optional[int] = None,
+    ae: edgedb.RelativeDuration,
+    af: typing.Optional[edgedb.RelativeDuration] = None,
+    ag: edgedb.DateDuration,
+    ah: typing.Optional[edgedb.DateDuration] = None,
+    ai: edgedb.ConfigMemory,
+    aj: typing.Optional[edgedb.ConfigMemory] = None,
+    ak: edgedb.Range[int],
+    al: typing.Optional[edgedb.Range[int]] = None,
+    am: edgedb.Range[int],
+    an: typing.Optional[edgedb.Range[int]] = None,
+    ao: edgedb.Range[float],
+    ap: typing.Optional[edgedb.Range[float]] = None,
+    aq: edgedb.Range[float],
+    ar: typing.Optional[edgedb.Range[float]] = None,
+    as_: edgedb.Range[datetime.datetime],
+    at: typing.Optional[edgedb.Range[datetime.datetime]] = None,
+    au: edgedb.Range[datetime.datetime],
+    av: typing.Optional[edgedb.Range[datetime.datetime]] = None,
+    aw: edgedb.Range[datetime.date],
+    ax: typing.Optional[edgedb.Range[datetime.date]] = None,
+    bc: typing.Sequence[float],
+    bd: typing.Optional[typing.Sequence[float]] = None,
+) -> MyQueryResult:
+    return await executor.query_single(
+        """\
+        create scalar type MyScalar extending int64;
+        create scalar type MyEnum extending enum<'This', 'That', '1', 'f. b', 'f-b'>;
+
+        select {
+        	a := <uuid>$a,
+        	b := <optional uuid>$b,
+        	c := <str>$c,
+        	d := <optional str>$d,
+        	e := <bytes>$e,
+        	f := <optional bytes>$f,
+        	g := <int16>$g,
+        	h := <optional int16>$h,
+        	i := <int32>$i,
+        	j := <optional int32>$j,
+        	k := <int64>$k,
+        	l := <optional int64>$l,
+        	m := <float32>$m,
+        	n := <optional float32>$n,
+        	o := <float64>$o,
+        	p := <optional float64>$p,
+        	q := <bool>$q,
+        	r := <optional bool>$r,
+        	s := <datetime>$s,
+        	t := <optional datetime>$t,
+        	u := <cal::local_datetime>$u,
+        	v := <optional cal::local_datetime>$v,
+        	w := <cal::local_date>$w,
+        	x := <optional cal::local_date>$x,
+        	y := <cal::local_time>$y,
+        	z := <optional cal::local_time>$z,
+        	aa := <duration>$aa,
+        	ab := <optional duration>$ab,
+        	ac := <bigint>$ac,
+        	ad := <optional bigint>$ad,
+        	ae := <cal::relative_duration>$ae,
+        	af := <optional cal::relative_duration>$af,
+        	ag := <cal::date_duration>$ag,
+        	ah := <optional cal::date_duration>$ah,
+        	ai := <cfg::memory>$ai,
+        	aj := <optional cfg::memory>$aj,
+        	ak := <range<int32>>$ak,
+        	al := <optional range<int32>>$al,
+        	am := <range<int64>>$am,
+        	an := <optional range<int64>>$an,
+        	ao := <range<float32>>$ao,
+        	ap := <optional range<float32>>$ap,
+        	aq := <range<float64>>$aq,
+        	ar := <optional range<float64>>$ar,
+        	as_ := <range<datetime>>$as_,
+        	at := <optional range<datetime>>$at,
+        	au := <range<cal::local_datetime>>$au,
+        	av := <optional range<cal::local_datetime>>$av,
+        	aw := <range<cal::local_date>>$aw,
+        	ax := <optional range<cal::local_date>>$ax,
+        	ay := <MyScalar>1,
+        	az := <optional MyScalar>{},
+        	ba := MyEnum.This,
+        	bb := <optional MyEnum>{},
+        	bc := <ext::pgvector::vector>$bc,
+        	bd := <optional ext::pgvector::vector>$bd,
+        }\
+        """,
+        a=a,
+        b=b,
+        c=c,
+        d=d,
+        e=e,
+        f=f,
+        g=g,
+        h=h,
+        i=i,
+        j=j,
+        k=k,
+        l=l,
+        m=m,
+        n=n,
+        o=o,
+        p=p,
+        q=q,
+        r=r,
+        s=s,
+        t=t,
+        u=u,
+        v=v,
+        w=w,
+        x=x,
+        y=y,
+        z=z,
+        aa=aa,
+        ab=ab,
+        ac=ac,
+        ad=ad,
+        ae=ae,
+        af=af,
+        ag=ag,
+        ah=ah,
+        ai=ai,
+        aj=aj,
+        ak=ak,
+        al=al,
+        am=am,
+        an=an,
+        ao=ao,
+        ap=ap,
+        aq=aq,
+        ar=ar,
+        as_=as_,
+        at=at,
+        au=au,
+        av=av,
+        aw=aw,
+        ax=ax,
+        bc=bc,
+        bd=bd,
+    )
diff --git a/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert b/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert
index c5a85ea8..fa79a94c 100644
--- a/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert
+++ b/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert
@@ -25,7 +25,6 @@ class MyEnum(enum.Enum):
 
 @dataclasses.dataclass
 class MyQueryResult:
-    id: uuid.UUID
     a: uuid.UUID
     b: typing.Optional[uuid.UUID]
     c: str
diff --git a/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert3 b/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert3
new file mode 120000
index 00000000..842e4d99
--- /dev/null
+++ b/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert3
@@ -0,0 +1,1 @@
+my_query_edgeql.py.assert5
\ No newline at end of file
diff --git a/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert5 b/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert5
new file mode 100644
index 00000000..c5a85ea8
--- /dev/null
+++ b/tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert5
@@ -0,0 +1,259 @@
+# AUTOGENERATED FROM 'parpkg/subpkg/my_query.edgeql' WITH:
+#     $ edgedb-py --target blocking --no-skip-pydantic-validation
+
+
+from __future__ import annotations
+import array
+import dataclasses
+import datetime
+import edgedb
+import enum
+import typing
+import uuid
+
+
+MyScalar = int
+
+
+class MyEnum(enum.Enum):
+    THIS = "This"
+    THAT = "That"
+    E_1 = "1"
+    F_B = "f. b"
+    F_B_1 = "f-b"
+
+
+@dataclasses.dataclass
+class MyQueryResult:
+    id: uuid.UUID
+    a: uuid.UUID
+    b: typing.Optional[uuid.UUID]
+    c: str
+    d: typing.Optional[str]
+    e: bytes
+    f: typing.Optional[bytes]
+    g: int
+    h: typing.Optional[int]
+    i: int
+    j: typing.Optional[int]
+    k: int
+    l: typing.Optional[int]
+    m: float
+    n: typing.Optional[float]
+    o: float
+    p: typing.Optional[float]
+    q: bool
+    r: typing.Optional[bool]
+    s: datetime.datetime
+    t: typing.Optional[datetime.datetime]
+    u: datetime.datetime
+    v: typing.Optional[datetime.datetime]
+    w: datetime.date
+    x: typing.Optional[datetime.date]
+    y: datetime.time
+    z: typing.Optional[datetime.time]
+    aa: datetime.timedelta
+    ab: typing.Optional[datetime.timedelta]
+    ac: int
+    ad: typing.Optional[int]
+    ae: edgedb.RelativeDuration
+    af: typing.Optional[edgedb.RelativeDuration]
+    ag: edgedb.DateDuration
+    ah: typing.Optional[edgedb.DateDuration]
+    ai: edgedb.ConfigMemory
+    aj: typing.Optional[edgedb.ConfigMemory]
+    ak: edgedb.Range[int]
+    al: typing.Optional[edgedb.Range[int]]
+    am: edgedb.Range[int]
+    an: typing.Optional[edgedb.Range[int]]
+    ao: edgedb.Range[float]
+    ap: typing.Optional[edgedb.Range[float]]
+    aq: edgedb.Range[float]
+    ar: typing.Optional[edgedb.Range[float]]
+    as_: edgedb.Range[datetime.datetime]
+    at: typing.Optional[edgedb.Range[datetime.datetime]]
+    au: edgedb.Range[datetime.datetime]
+    av: typing.Optional[edgedb.Range[datetime.datetime]]
+    aw: edgedb.Range[datetime.date]
+    ax: typing.Optional[edgedb.Range[datetime.date]]
+    ay: MyScalar
+    az: typing.Optional[MyScalar]
+    ba: MyEnum
+    bb: typing.Optional[MyEnum]
+    bc: array.array
+    bd: typing.Optional[array.array]
+
+
+def my_query(
+    executor: edgedb.Executor,
+    *,
+    a: uuid.UUID,
+    b: typing.Optional[uuid.UUID] = None,
+    c: str,
+    d: typing.Optional[str] = None,
+    e: bytes,
+    f: typing.Optional[bytes] = None,
+    g: int,
+    h: typing.Optional[int] = None,
+    i: int,
+    j: typing.Optional[int] = None,
+    k: int,
+    l: typing.Optional[int] = None,
+    m: float,
+    n: typing.Optional[float] = None,
+    o: float,
+    p: typing.Optional[float] = None,
+    q: bool,
+    r: typing.Optional[bool] = None,
+    s: datetime.datetime,
+    t: typing.Optional[datetime.datetime] = None,
+    u: datetime.datetime,
+    v: typing.Optional[datetime.datetime] = None,
+    w: datetime.date,
+    x: typing.Optional[datetime.date] = None,
+    y: datetime.time,
+    z: typing.Optional[datetime.time] = None,
+    aa: datetime.timedelta,
+    ab: typing.Optional[datetime.timedelta] = None,
+    ac: int,
+    ad: typing.Optional[int] = None,
+    ae: edgedb.RelativeDuration,
+    af: typing.Optional[edgedb.RelativeDuration] = None,
+    ag: edgedb.DateDuration,
+    ah: typing.Optional[edgedb.DateDuration] = None,
+    ai: edgedb.ConfigMemory,
+    aj: typing.Optional[edgedb.ConfigMemory] = None,
+    ak: edgedb.Range[int],
+    al: typing.Optional[edgedb.Range[int]] = None,
+    am: edgedb.Range[int],
+    an: typing.Optional[edgedb.Range[int]] = None,
+    ao: edgedb.Range[float],
+    ap: typing.Optional[edgedb.Range[float]] = None,
+    aq: edgedb.Range[float],
+    ar: typing.Optional[edgedb.Range[float]] = None,
+    as_: edgedb.Range[datetime.datetime],
+    at: typing.Optional[edgedb.Range[datetime.datetime]] = None,
+    au: edgedb.Range[datetime.datetime],
+    av: typing.Optional[edgedb.Range[datetime.datetime]] = None,
+    aw: edgedb.Range[datetime.date],
+    ax: typing.Optional[edgedb.Range[datetime.date]] = None,
+    bc: typing.Sequence[float],
+    bd: typing.Optional[typing.Sequence[float]] = None,
+) -> MyQueryResult:
+    return executor.query_single(
+        """\
+        create scalar type MyScalar extending int64;
+        create scalar type MyEnum extending enum<'This', 'That', '1', 'f. b', 'f-b'>;
+
+        select {
+        	a := <uuid>$a,
+        	b := <optional uuid>$b,
+        	c := <str>$c,
+        	d := <optional str>$d,
+        	e := <bytes>$e,
+        	f := <optional bytes>$f,
+        	g := <int16>$g,
+        	h := <optional int16>$h,
+        	i := <int32>$i,
+        	j := <optional int32>$j,
+        	k := <int64>$k,
+        	l := <optional int64>$l,
+        	m := <float32>$m,
+        	n := <optional float32>$n,
+        	o := <float64>$o,
+        	p := <optional float64>$p,
+        	q := <bool>$q,
+        	r := <optional bool>$r,
+        	s := <datetime>$s,
+        	t := <optional datetime>$t,
+        	u := <cal::local_datetime>$u,
+        	v := <optional cal::local_datetime>$v,
+        	w := <cal::local_date>$w,
+        	x := <optional cal::local_date>$x,
+        	y := <cal::local_time>$y,
+        	z := <optional cal::local_time>$z,
+        	aa := <duration>$aa,
+        	ab := <optional duration>$ab,
+        	ac := <bigint>$ac,
+        	ad := <optional bigint>$ad,
+        	ae := <cal::relative_duration>$ae,
+        	af := <optional cal::relative_duration>$af,
+        	ag := <cal::date_duration>$ag,
+        	ah := <optional cal::date_duration>$ah,
+        	ai := <cfg::memory>$ai,
+        	aj := <optional cfg::memory>$aj,
+        	ak := <range<int32>>$ak,
+        	al := <optional range<int32>>$al,
+        	am := <range<int64>>$am,
+        	an := <optional range<int64>>$an,
+        	ao := <range<float32>>$ao,
+        	ap := <optional range<float32>>$ap,
+        	aq := <range<float64>>$aq,
+        	ar := <optional range<float64>>$ar,
+        	as_ := <range<datetime>>$as_,
+        	at := <optional range<datetime>>$at,
+        	au := <range<cal::local_datetime>>$au,
+        	av := <optional range<cal::local_datetime>>$av,
+        	aw := <range<cal::local_date>>$aw,
+        	ax := <optional range<cal::local_date>>$ax,
+        	ay := <MyScalar>1,
+        	az := <optional MyScalar>{},
+        	ba := MyEnum.This,
+        	bb := <optional MyEnum>{},
+        	bc := <ext::pgvector::vector>$bc,
+        	bd := <optional ext::pgvector::vector>$bd,
+        }\
+        """,
+        a=a,
+        b=b,
+        c=c,
+        d=d,
+        e=e,
+        f=f,
+        g=g,
+        h=h,
+        i=i,
+        j=j,
+        k=k,
+        l=l,
+        m=m,
+        n=n,
+        o=o,
+        p=p,
+        q=q,
+        r=r,
+        s=s,
+        t=t,
+        u=u,
+        v=v,
+        w=w,
+        x=x,
+        y=y,
+        z=z,
+        aa=aa,
+        ab=ab,
+        ac=ac,
+        ad=ad,
+        ae=ae,
+        af=af,
+        ag=ag,
+        ah=ah,
+        ai=ai,
+        aj=aj,
+        ak=ak,
+        al=al,
+        am=am,
+        an=an,
+        ao=ao,
+        ap=ap,
+        aq=aq,
+        ar=ar,
+        as_=as_,
+        at=at,
+        au=au,
+        av=av,
+        aw=aw,
+        ax=ax,
+        bc=bc,
+        bd=bd,
+    )
diff --git a/tests/codegen/test-project2/scalar/custom_vector_input_async_edgeql.py.assert b/tests/codegen/test-project2/scalar/custom_vector_input_async_edgeql.py.assert
index a16beb5e..39272b4f 100644
--- a/tests/codegen/test-project2/scalar/custom_vector_input_async_edgeql.py.assert
+++ b/tests/codegen/test-project2/scalar/custom_vector_input_async_edgeql.py.assert
@@ -7,13 +7,13 @@ import edgedb
 import typing
 
 
-Input = typing.Sequence[float]
+V3 = typing.Sequence[float]
 
 
 async def custom_vector_input(
     executor: edgedb.AsyncIOExecutor,
     *,
-    input: typing.Optional[Input] = None,
+    input: typing.Optional[V3] = None,
 ) -> typing.Optional[int]:
     return await executor.query_single(
         """\
diff --git a/tests/codegen/test-project2/scalar/custom_vector_input_async_edgeql.py.assert4 b/tests/codegen/test-project2/scalar/custom_vector_input_async_edgeql.py.assert3
similarity index 84%
rename from tests/codegen/test-project2/scalar/custom_vector_input_async_edgeql.py.assert4
rename to tests/codegen/test-project2/scalar/custom_vector_input_async_edgeql.py.assert3
index 39272b4f..a16beb5e 100644
--- a/tests/codegen/test-project2/scalar/custom_vector_input_async_edgeql.py.assert4
+++ b/tests/codegen/test-project2/scalar/custom_vector_input_async_edgeql.py.assert3
@@ -7,13 +7,13 @@ import edgedb
 import typing
 
 
-V3 = typing.Sequence[float]
+Input = typing.Sequence[float]
 
 
 async def custom_vector_input(
     executor: edgedb.AsyncIOExecutor,
     *,
-    input: typing.Optional[V3] = None,
+    input: typing.Optional[Input] = None,
 ) -> typing.Optional[int]:
     return await executor.query_single(
         """\
diff --git a/tests/codegen/test-project2/scalar/custom_vector_input_edgeql.py.assert b/tests/codegen/test-project2/scalar/custom_vector_input_edgeql.py.assert
index 4cf4659a..ed2947c0 100644
--- a/tests/codegen/test-project2/scalar/custom_vector_input_edgeql.py.assert
+++ b/tests/codegen/test-project2/scalar/custom_vector_input_edgeql.py.assert
@@ -7,13 +7,13 @@ import edgedb
 import typing
 
 
-Input = typing.Sequence[float]
+V3 = typing.Sequence[float]
 
 
 def custom_vector_input(
     executor: edgedb.Executor,
     *,
-    input: typing.Optional[Input] = None,
+    input: typing.Optional[V3] = None,
 ) -> typing.Optional[int]:
     return executor.query_single(
         """\
diff --git a/tests/codegen/test-project2/scalar/custom_vector_input_edgeql.py.assert4 b/tests/codegen/test-project2/scalar/custom_vector_input_edgeql.py.assert3
similarity index 85%
rename from tests/codegen/test-project2/scalar/custom_vector_input_edgeql.py.assert4
rename to tests/codegen/test-project2/scalar/custom_vector_input_edgeql.py.assert3
index ed2947c0..4cf4659a 100644
--- a/tests/codegen/test-project2/scalar/custom_vector_input_edgeql.py.assert4
+++ b/tests/codegen/test-project2/scalar/custom_vector_input_edgeql.py.assert3
@@ -7,13 +7,13 @@ import edgedb
 import typing
 
 
-V3 = typing.Sequence[float]
+Input = typing.Sequence[float]
 
 
 def custom_vector_input(
     executor: edgedb.Executor,
     *,
-    input: typing.Optional[V3] = None,
+    input: typing.Optional[Input] = None,
 ) -> typing.Optional[int]:
     return executor.query_single(
         """\
diff --git a/tests/test_codegen.py b/tests/test_codegen.py
index 831a3530..801b00ae 100644
--- a/tests/test_codegen.py
+++ b/tests/test_codegen.py
@@ -25,11 +25,9 @@
 import tempfile
 
 from edgedb import _testbase as tb
-import unittest
 
 
-# Use ".assert" for EdgeDB 3.x and lower
-ASSERT_SUFFIX = os.environ.get("EDGEDB_TEST_CODEGEN_ASSERT_SUFFIX", ".assert4")
+ASSERT_SUFFIX = os.environ.get("EDGEDB_TEST_CODEGEN_ASSERT_SUFFIX", ".assert")
 
 
 class TestCodegen(tb.AsyncQueryTestCase):
@@ -43,11 +41,6 @@ class TestCodegen(tb.AsyncQueryTestCase):
         drop extension pgvector;
     '''
 
-    @unittest.skip('''
-        The codegen seems to be broken w.r.t. expectations on whether `id` is
-        supposed to appear in `MyQueryResult` in the
-        `generated_async_edgeql.py.assert4`.
-    ''')
     async def test_codegen(self):
         env = os.environ.copy()
         env.update(