From 62358198a993860ef002f150b2b13fba4d636401 Mon Sep 17 00:00:00 2001 From: Devin Stein Date: Wed, 22 Jan 2025 21:31:36 -0800 Subject: [PATCH 1/2] chore: bump README to 1.3.0 --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 5d53009..c9f2531 100644 --- a/README.md +++ b/README.md @@ -5,8 +5,8 @@ version: "2" plugins: - name: py wasm: - url: https://downloads.sqlc.dev/plugin/sqlc-gen-python_1.2.0.wasm - sha256: a6c5d174c407007c3717eea36ff0882744346e6ba991f92f71d6ab2895204c0e + url: https://downloads.sqlc.dev/plugin/sqlc-gen-python_1.3.0.wasm + sha256: fbedae96b5ecae2380a70fb5b925fd4bff58a6cfb1f3140375d098fbab7b3a3c sql: - schema: "schema.sql" queries: "query.sql" From 277ec542b98d9167ffb3643bbcc371ab954b6c41 Mon Sep 17 00:00:00 2001 From: Dave Hein Date: Tue, 16 Dec 2025 18:48:32 -0600 Subject: [PATCH 2/2] feat: add pgvector support with typed numpy arrays Add support for PostgreSQL vector type with proper numpy typing: - Map 'vector' type to NDArray[numpy.float32] - Auto-import numpy and numpy.typing.NDArray - Add test case for vector type generation - Add .DS_Store to .gitignore This enables type-safe pgvector usage in generated Python code. --- .gitignore | 3 + .../emit_numpy_array/python/models.py | 13 ++++ .../testdata/emit_numpy_array/python/query.py | 68 +++++++++++++++++++ .../testdata/emit_numpy_array/query.sql | 6 ++ .../testdata/emit_numpy_array/schema.sql | 4 ++ .../testdata/emit_numpy_array/sqlc.yaml | 17 +++++ internal/imports.go | 4 ++ internal/postgresql_type.go | 2 + 8 files changed, 117 insertions(+) create mode 100644 internal/endtoend/testdata/emit_numpy_array/python/models.py create mode 100644 internal/endtoend/testdata/emit_numpy_array/python/query.py create mode 100644 internal/endtoend/testdata/emit_numpy_array/query.sql create mode 100644 internal/endtoend/testdata/emit_numpy_array/schema.sql create mode 100644 internal/endtoend/testdata/emit_numpy_array/sqlc.yaml diff --git a/.gitignore b/.gitignore index 60b4f3d..0d3523e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,8 @@ bin +# macOS +.DS_Store + # Devenv .envrc .direnv diff --git a/internal/endtoend/testdata/emit_numpy_array/python/models.py b/internal/endtoend/testdata/emit_numpy_array/python/models.py new file mode 100644 index 0000000..058ad94 --- /dev/null +++ b/internal/endtoend/testdata/emit_numpy_array/python/models.py @@ -0,0 +1,13 @@ +# Code generated by sqlc. DO NOT EDIT. +# versions: +# sqlc v1.30.0 +import dataclasses +import numpy +from numpy.typing import NDArray +from typing import Optional + + +@dataclasses.dataclass() +class Item: + id: int + embedding: Optional[NDArray[numpy.float32]] diff --git a/internal/endtoend/testdata/emit_numpy_array/python/query.py b/internal/endtoend/testdata/emit_numpy_array/python/query.py new file mode 100644 index 0000000..cdcdd02 --- /dev/null +++ b/internal/endtoend/testdata/emit_numpy_array/python/query.py @@ -0,0 +1,68 @@ +# Code generated by sqlc. DO NOT EDIT. +# versions: +# sqlc v1.30.0 +# source: query.sql +import numpy +from numpy.typing import NDArray +from typing import Optional + +import sqlalchemy +import sqlalchemy.ext.asyncio + +from python import models + + +CREATE_ITEM = """-- name: create_item \\:one +INSERT INTO items (embedding) VALUES (:p1) RETURNING id, embedding +""" + + +GET_ITEM = """-- name: get_item \\:one +SELECT id, embedding FROM items WHERE id = :p1 +""" + + +class Querier: + def __init__(self, conn: sqlalchemy.engine.Connection): + self._conn = conn + + def create_item(self, *, embedding: Optional[NDArray[numpy.float32]]) -> Optional[models.Item]: + row = self._conn.execute(sqlalchemy.text(CREATE_ITEM), {"p1": embedding}).first() + if row is None: + return None + return models.Item( + id=row[0], + embedding=row[1], + ) + + def get_item(self, *, id: int) -> Optional[models.Item]: + row = self._conn.execute(sqlalchemy.text(GET_ITEM), {"p1": id}).first() + if row is None: + return None + return models.Item( + id=row[0], + embedding=row[1], + ) + + +class AsyncQuerier: + def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection): + self._conn = conn + + async def create_item(self, *, embedding: Optional[NDArray[numpy.float32]]) -> Optional[models.Item]: + row = (await self._conn.execute(sqlalchemy.text(CREATE_ITEM), {"p1": embedding})).first() + if row is None: + return None + return models.Item( + id=row[0], + embedding=row[1], + ) + + async def get_item(self, *, id: int) -> Optional[models.Item]: + row = (await self._conn.execute(sqlalchemy.text(GET_ITEM), {"p1": id})).first() + if row is None: + return None + return models.Item( + id=row[0], + embedding=row[1], + ) diff --git a/internal/endtoend/testdata/emit_numpy_array/query.sql b/internal/endtoend/testdata/emit_numpy_array/query.sql new file mode 100644 index 0000000..e1bfc30 --- /dev/null +++ b/internal/endtoend/testdata/emit_numpy_array/query.sql @@ -0,0 +1,6 @@ +-- name: GetItem :one +SELECT * FROM items WHERE id = $1; + + +-- name: CreateItem :one +INSERT INTO items (embedding) VALUES ($1) RETURNING *; \ No newline at end of file diff --git a/internal/endtoend/testdata/emit_numpy_array/schema.sql b/internal/endtoend/testdata/emit_numpy_array/schema.sql new file mode 100644 index 0000000..e347cd8 --- /dev/null +++ b/internal/endtoend/testdata/emit_numpy_array/schema.sql @@ -0,0 +1,4 @@ +CREATE TABLE items ( + id SERIAL PRIMARY KEY, + embedding vector(3) +); diff --git a/internal/endtoend/testdata/emit_numpy_array/sqlc.yaml b/internal/endtoend/testdata/emit_numpy_array/sqlc.yaml new file mode 100644 index 0000000..a84ffdf --- /dev/null +++ b/internal/endtoend/testdata/emit_numpy_array/sqlc.yaml @@ -0,0 +1,17 @@ +version: '2' +plugins: +- name: py + wasm: + url: file://../../../../bin/sqlc-gen-python.wasm + sha256: "839af1f07c31644548192fc095569e62f1511d72c1c30c1a958ddc9c9429edbc" +sql: +- schema: schema.sql + queries: query.sql + engine: postgresql + codegen: + - plugin: py + out: python + options: + package: python + emit_sync_querier: true + emit_async_querier: true diff --git a/internal/imports.go b/internal/imports.go index b88c58c..ff2f731 100644 --- a/internal/imports.go +++ b/internal/imports.go @@ -272,5 +272,9 @@ func stdImports(uses func(name string) bool) map[string]importSpec { if uses("Any") { std["typing.Any"] = importSpec{Module: "typing", Name: "Any"} } + if uses("NDArray[numpy.float32]") { + std["numpy"] = importSpec{Module: "numpy"} + std["numpy.typing.NDArray"] = importSpec{Module: "numpy.typing", Name: "NDArray"} + } return std } diff --git a/internal/postgresql_type.go b/internal/postgresql_type.go index 3d0891b..019b097 100644 --- a/internal/postgresql_type.go +++ b/internal/postgresql_type.go @@ -42,6 +42,8 @@ func postgresType(req *plugin.GenerateRequest, col *plugin.Column) string { return "str" case "ltree", "lquery", "ltxtquery": return "str" + case "vector": + return "NDArray[numpy.float32]" default: for _, schema := range req.Catalog.Schemas { if schema.Name == "pg_catalog" || schema.Name == "information_schema" {