Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix codegen for pgvector #447

Merged
merged 4 commits into from
Sep 14, 2023
Merged

Fix codegen for pgvector #447

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
cfa3db6
Fix codegen for pgvector
fantix Jun 26, 2023
117efe5
Merge branch 'master' into codegen-pgvector
fantix Sep 8, 2023
c7d66ae
CRF: use Sequence[float] for vector input
fantix Sep 8, 2023
c897eae
Merge branch 'master' into codegen-pgvector
fantix Sep 14, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions edgedb/codegen/generator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@
"cal::relative_duration": "edgedb.RelativeDuration",
"cal::date_duration": "edgedb.DateDuration",
"cfg::memory": "edgedb.ConfigMemory",
"ext::pgvector::vector": "array.array",
}

TYPE_IMPORTS = {
Expand All @@ -77,6 +78,7 @@
"cal::local_date": "datetime",
"cal::local_time": "datetime",
"cal::local_datetime": "datetime",
"ext::pgvector::vector": "array",
}

PYDANTIC_MIXIN = """\
Expand Down
9 changes: 9 additions & 0 deletions tests/codegen/test-project2/generated_async_edgeql.py.assert
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@


from __future__ import annotations
import array
import dataclasses
import datetime
import edgedb
Expand Down Expand Up @@ -113,6 +114,8 @@ class MyQueryResult:
az: MyScalar | None
ba: MyEnum
bb: MyEnum | None
bc: array.array
fantix marked this conversation as resolved.
Show resolved Hide resolved
bd: array.array | None


@dataclasses.dataclass
Expand Down Expand Up @@ -217,6 +220,8 @@ async def my_query(
av: edgedb.Range[datetime.datetime] | None = None,
aw: edgedb.Range[datetime.date],
ax: edgedb.Range[datetime.date] | None = None,
bc: array.array,
bd: array.array | None = None,
) -> MyQueryResult:
return await executor.query_single(
"""\
Expand Down Expand Up @@ -278,6 +283,8 @@ async def my_query(
az := <optional MyScalar>{},
ba := MyEnum.This,
bb := <optional MyEnum>{},
bc := <ext::pgvector::vector>$bc,
bd := <optional ext::pgvector::vector>$bd,
}\
""",
a=a,
Expand Down Expand Up @@ -330,6 +337,8 @@ async def my_query(
av=av,
aw=aw,
ax=ax,
bc=bc,
bd=bd,
)


Expand Down
2 changes: 2 additions & 0 deletions tests/codegen/test-project2/parpkg/subpkg/my_query.edgeql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,4 +56,6 @@ select {
az := <optional MyScalar>{},
ba := MyEnum.This,
bb := <optional MyEnum>{},
bc := <ext::pgvector::vector>$bc,
bd := <optional ext::pgvector::vector>$bd,
}
9 changes: 9 additions & 0 deletions tests/codegen/test-project2/parpkg/subpkg/my_query_async_edgeql.py.assert
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@


from __future__ import annotations
import array
import dataclasses
import datetime
import edgedb
Expand Down Expand Up @@ -88,6 +89,8 @@ class MyQueryResult(NoPydanticValidation):
az: typing.Optional[MyScalar]
ba: MyEnum
bb: typing.Optional[MyEnum]
bc: array.array
bd: typing.Optional[array.array]


async def my_query(
Expand Down Expand Up @@ -143,6 +146,8 @@ async def my_query(
av: typing.Optional[edgedb.Range[datetime.datetime]] = None,
aw: edgedb.Range[datetime.date],
ax: typing.Optional[edgedb.Range[datetime.date]] = None,
bc: array.array,
bd: typing.Optional[array.array] = None,
) -> MyQueryResult:
return await executor.query_single(
"""\
Expand Down Expand Up @@ -204,6 +209,8 @@ async def my_query(
az := <optional MyScalar>{},
ba := MyEnum.This,
bb := <optional MyEnum>{},
bc := <ext::pgvector::vector>$bc,
bd := <optional ext::pgvector::vector>$bd,
}\
""",
a=a,
Expand Down Expand Up @@ -256,4 +263,6 @@ async def my_query(
av=av,
aw=aw,
ax=ax,
bc=bc,
bd=bd,
)
9 changes: 9 additions & 0 deletions tests/codegen/test-project2/parpkg/subpkg/my_query_edgeql.py.assert
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@


from __future__ import annotations
import array
import dataclasses
import datetime
import edgedb
Expand Down Expand Up @@ -79,6 +80,8 @@ class MyQueryResult:
az: typing.Optional[MyScalar]
ba: MyEnum
bb: typing.Optional[MyEnum]
bc: array.array
bd: typing.Optional[array.array]


def my_query(
Expand Down Expand Up @@ -134,6 +137,8 @@ def my_query(
av: typing.Optional[edgedb.Range[datetime.datetime]] = None,
aw: edgedb.Range[datetime.date],
ax: typing.Optional[edgedb.Range[datetime.date]] = None,
bc: array.array,
bd: typing.Optional[array.array] = None,
) -> MyQueryResult:
return executor.query_single(
"""\
Expand Down Expand Up @@ -195,6 +200,8 @@ def my_query(
az := <optional MyScalar>{},
ba := MyEnum.This,
bb := <optional MyEnum>{},
bc := <ext::pgvector::vector>$bc,
bd := <optional ext::pgvector::vector>$bd,
}\
""",
a=a,
Expand Down Expand Up @@ -247,4 +254,6 @@ def my_query(
av=av,
aw=aw,
ax=ax,
bc=bc,
bd=bd,
)
9 changes: 9 additions & 0 deletions tests/test_codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,14 @@


class TestCodegen(tb.AsyncQueryTestCase):
SETUP = '''
create extension pgvector;
'''

TEARDOWN = '''
drop extension pgvector;
'''

async def test_codegen(self):
env = os.environ.copy()
env.update(
Expand All @@ -36,6 +44,7 @@ async def test_codegen(self):
for k, v in self.get_connect_args().items()
}
)
env["EDGEDB_DATABASE"] = self.get_database_name()
container = pathlib.Path(__file__).absolute().parent / "codegen"
with tempfile.TemporaryDirectory() as td:
td_path = pathlib.Path(td)
Expand Down