Skip to content

Commit 3754349

Browse files
kyleconroykyleconroy
authored
feat(codegen): Add support for pgvector and pgvector-go (#2888)
* feat(codegen): Add support for pgvector-go * test: Update ddl_tests to use FindTests
1 parent bd9919d commit 3754349

File tree

12 files changed

+166
-23
lines changed

12 files changed

+166
-23
lines changed

‎internal/codegen/golang/imports.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,10 @@ func buildImports(settings *plugin.Settings, options *opts, queries []Query, use
219219
if uses("uuid.NullUUID") && !overrideNullUUID {
220220
pkg[ImportSpec{Path: "github.com/google/uuid"}] = struct{}{}
221221
}
222+
_, overrideVector := overrideTypes["pgvector.Vector"]
223+
if uses("pgvector.Vector") && !overrideVector {
224+
pkg[ImportSpec{Path: "github.com/pgvector/pgvector-go"}] = struct{}{}
225+
}
222226

223227
// Custom imports
224228
for _, o := range settings.Overrides {

‎internal/codegen/golang/postgresql_type.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -537,6 +537,15 @@ func postgresType(req *plugin.CodeGenRequest, options *opts, col *plugin.Column)
537537
return "pgtype.Polygon"
538538
}
539539

540+
case "vector":
541+
if driver == SQLDriverPGXV5 {
542+
if emitPointersForNull {
543+
return "*pgvector.Vector"
544+
} else {
545+
return "pgvector.Vector"
546+
}
547+
}
548+
540549
case "void":
541550
// A void value can only be scanned into an empty interface.
542551
return "interface{}"

‎internal/endtoend/ddl_test.go

Lines changed: 10 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@ package main
33
import (
44
"context"
55
"fmt"
6-
"io/fs"
76
"os"
87
"path/filepath"
8+
"slices"
99
"strings"
1010
"testing"
1111

@@ -32,28 +32,20 @@ func TestValidSchema(t *testing.T) {
3232
t.Fatal(err)
3333
}
3434

35-
files := []string{}
35+
for _, replay := range FindTests(t, "testdata", "managed-db") {
36+
replay := replay // https://golang.org/doc/faq#closures_and_goroutines
3637

37-
// Find all tests that do not have a stderr.txt file
38-
err = filepath.Walk("testdata", func(path string, info fs.FileInfo, err error) error {
39-
if err != nil {
40-
return err
38+
if len(replay.Stderr) > 0 {
39+
continue
4140
}
42-
if filepath.Base(path) == "sqlc.json" || filepath.Base(path) == "sqlc.yaml" || filepath.Base(path) == "sqlc.yml" {
43-
stderr := filepath.Join(filepath.Dir(path), "stderr.txt")
44-
if _, err := os.Stat(stderr); !os.IsNotExist(err) {
45-
return nil
41+
42+
if replay.Exec != nil {
43+
if !slices.Contains(replay.Exec.Contexts, "managed-db") {
44+
continue
4645
}
47-
files = append(files, path)
4846
}
49-
return nil
50-
})
51-
if err != nil {
52-
t.Fatal(err)
53-
}
5447

55-
for _, file := range files {
56-
file := file // https://golang.org/doc/faq#closures_and_goroutines
48+
file := filepath.Join(replay.Path, replay.ConfigName)
5749
rd, err := os.Open(file)
5850
if err != nil {
5951
t.Fatal(err)

‎internal/endtoend/testdata/go.mod

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ require (
1010
github.com/jackc/pgconn v1.5.1-0.20200601181101-fa742c524853
1111
github.com/jackc/pgtype v1.6.2
1212
github.com/jackc/pgx/v4 v4.6.1-0.20200606145419-4e5062306904
13-
github.com/jackc/pgx/v5 v5.0.1
13+
github.com/jackc/pgx/v5 v5.4.3
1414
github.com/lib/pq v1.9.0
1515
github.com/sqlc-dev/pqtype v0.2.0
1616
github.com/sqlc-dev/sqlc-testdata v1.0.0
@@ -24,11 +24,12 @@ require (
2424
github.com/jackc/pgio v1.0.0 // indirect
2525
github.com/jackc/pgpassfile v1.0.0 // indirect
2626
github.com/jackc/pgproto3/v2 v2.0.1 // indirect
27-
github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b // indirect
27+
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
28+
github.com/pgvector/pgvector-go v0.1.1 // indirect
2829
github.com/volatiletech/inflect v0.0.1 // indirect
2930
github.com/volatiletech/randomize v0.0.1 // indirect
3031
github.com/volatiletech/strmangle v0.0.1 // indirect
31-
golang.org/x/crypto v0.0.0-20220829220503-c86fa9a7ed90 // indirect
32-
golang.org/x/text v0.3.7 // indirect
33-
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 // indirect
32+
golang.org/x/crypto v0.9.0 // indirect
33+
golang.org/x/text v0.9.0 // indirect
34+
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
3435
)

‎internal/endtoend/testdata/go.sum

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,8 @@ github.com/jackc/pgproto3/v2 v2.0.1/go.mod h1:WfJCnwN3HIg9Ish/j3sgWXnAfK8A9Y0bwX
4747
github.com/jackc/pgservicefile v0.0.0-20200307190119-3430c5407db8/go.mod h1:vsD4gTJCa9TptPL8sPkXrLZ+hDuNrZCnj29CQpr4X1E=
4848
github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b h1:C8S2+VttkHFdOOCXJe+YGfa4vHYwlt4Zx+IVXQ97jYg=
4949
github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b/go.mod h1:vsD4gTJCa9TptPL8sPkXrLZ+hDuNrZCnj29CQpr4X1E=
50+
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk=
51+
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
5052
github.com/jackc/pgtype v0.0.0-20190421001408-4ed0de4755e0/go.mod h1:hdSHsc1V01CGwFsrv11mJRHWJ6aifDLfdV3aVjFF0zg=
5153
github.com/jackc/pgtype v0.0.0-20190824184912-ab885b375b90/go.mod h1:KcahbBH1nCMSo2DXpzsoWOAfFkdEtEJpPbVLq8eE+mc=
5254
github.com/jackc/pgtype v0.0.0-20190828014616-a8802b16cc59/go.mod h1:MWlu30kVJrUS8lot6TQqcg7mtthZ9T0EoIBFiJcmcyw=
@@ -64,6 +66,10 @@ github.com/jackc/pgx/v4 v4.6.1-0.20200606145419-4e5062306904 h1:SdGWuGg+Cpxq6Z+A
6466
github.com/jackc/pgx/v4 v4.6.1-0.20200606145419-4e5062306904/go.mod h1:ZDaNWkt9sW1JMiNn0kdYBaLelIhw7Pg4qd+Vk6tw7Hg=
6567
github.com/jackc/pgx/v5 v5.0.1 h1:JZu9othr7l8so2JMDAGeDUMXqERAuZpovyfl4H50tdg=
6668
github.com/jackc/pgx/v5 v5.0.1/go.mod h1:JBbvW3Hdw77jKl9uJrEDATUZIFM2VFPzRq4RWIhkF4o=
69+
github.com/jackc/pgx/v5 v5.3.1 h1:Fcr8QJ1ZeLi5zsPZqQeUZhNhxfkkKBOgJuYkJHoBOtU=
70+
github.com/jackc/pgx/v5 v5.3.1/go.mod h1:t3JDKnCBlYIc0ewLF0Q7B8MXmoIaBOZj/ic7iHozM/8=
71+
github.com/jackc/pgx/v5 v5.4.3 h1:cxFyXhxlvAifxnkKKdlxv8XqUf59tDlYjnV5YYfsJJY=
72+
github.com/jackc/pgx/v5 v5.4.3/go.mod h1:Ig06C2Vu0t5qXC60W8sqIthScaEnFvojjj9dSljmHRA=
6773
github.com/jackc/puddle v0.0.0-20190413234325-e4ced69a3a2b/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
6874
github.com/jackc/puddle v0.0.0-20190608224051-11cab39313c9/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
6975
github.com/jackc/puddle v1.1.0/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
@@ -89,6 +95,8 @@ github.com/mattn/go-isatty v0.0.7/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hd
8995
github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
9096
github.com/mattn/go-isatty v0.0.9/go.mod h1:YNRxwqDuOph6SZLI9vUUz6OYw3QyUt7WiY2yME+cCiQ=
9197
github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
98+
github.com/pgvector/pgvector-go v0.1.1 h1:kqJigGctFnlWvskUiYIvJRNwUtQl/aMSUZVs0YWQe+g=
99+
github.com/pgvector/pgvector-go v0.1.1/go.mod h1:wLJgD/ODkdtd2LJK4l6evHXTuG+8PxymYAVomKHOWac=
92100
github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
93101
github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
94102
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
@@ -114,6 +122,7 @@ github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXf
114122
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
115123
github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
116124
github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
125+
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
117126
github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
118127
github.com/volatiletech/inflect v0.0.1 h1:2a6FcMQyhmPZcLa+uet3VJ8gLn/9svWhJxJYwvE8KsU=
119128
github.com/volatiletech/inflect v0.0.1/go.mod h1:IBti31tG6phkHitLlr5j7shC5SOo//x0AjDzaJU1PLA=
@@ -140,6 +149,9 @@ golang.org/x/crypto v0.0.0-20190911031432-227b76d455e7/go.mod h1:yigFU9vqHzYiE8U
140149
golang.org/x/crypto v0.0.0-20200323165209-0ec3e9974c59/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
141150
golang.org/x/crypto v0.0.0-20220829220503-c86fa9a7ed90 h1:Y/gsMcFOcR+6S6f3YeMKl5g+dZMEWqcz5Czj/GWYbkM=
142151
golang.org/x/crypto v0.0.0-20220829220503-c86fa9a7ed90/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
152+
golang.org/x/crypto v0.6.0/go.mod h1:OFC/31mSvZgRz0V1QTNCzfAI1aIRzbiufJtkMIlEp58=
153+
golang.org/x/crypto v0.9.0 h1:LF6fAI+IutBocDJ2OT0Q1g8plpYljMZ4+lty+dsqw3g=
154+
golang.org/x/crypto v0.9.0/go.mod h1:yrmDGqONDYtNj3tH8X9dzUun2m2lzPa9ngI6/RUPGR0=
143155
golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
144156
golang.org/x/mod v0.0.0-20190513183733-4bf6d317e70e/go.mod h1:mXi4GBBbnImb6dmsKGUJ2LatrhH/nqhxcFungHvyanc=
145157
golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
@@ -161,6 +173,10 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
161173
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
162174
golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
163175
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
176+
golang.org/x/text v0.7.0 h1:4BRB4x83lYWy72KwLD/qYDuTu7q9PjSagHvijDw7cLo=
177+
golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
178+
golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
179+
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
164180
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
165181
golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
166182
golang.org/x/tools v0.0.0-20190425163242-31fd60d6bfdc/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
@@ -173,12 +189,15 @@ golang.org/x/xerrors v0.0.0-20190513163551-3ee3066db522/go.mod h1:I/5z698sn9Ka8T
173189
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
174190
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
175191
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
192+
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=
193+
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
176194
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
177195
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
178196
gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
179197
gopkg.in/guregu/null.v4 v4.0.0 h1:1Wm3S1WEA2I26Kq+6vcW+w0gcDo44YKYD7YIEJNHDjg=
180198
gopkg.in/guregu/null.v4 v4.0.0/go.mod h1:YoQhUrADuG3i9WqesrCmpNRwm1ypAgSHYqoOcTu/JrI=
181199
gopkg.in/inconshreveable/log15.v2 v2.0.0-20180818164646-67afb5ed74ec/go.mod h1:aPpfJ7XW+gOuirDoZ8gHhLh3kZ1B08FtV2bbmy7Jv3s=
182200
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
201+
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
183202
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
184203
honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=

‎internal/endtoend/testdata/pg_vector/postgresql/pgx/exec.json

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
{
2+
"contexts": ["base"]
3+
}

‎internal/endtoend/testdata/pg_vector/postgresql/pgx/go/db.go

Lines changed: 32 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎internal/endtoend/testdata/pg_vector/postgresql/pgx/go/models.go

Lines changed: 14 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎internal/endtoend/testdata/pg_vector/postgresql/pgx/go/query.sql.go

Lines changed: 48 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎internal/endtoend/testdata/pg_vector/postgresql/pgx/query.sql

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
-- name: InsertVector :exec
2+
INSERT INTO items (embedding) VALUES ($1);
3+
4+
-- name: NearestNeighbor :many
5+
SELECT *
6+
FROM items
7+
ORDER BY embedding <-> $1
8+
LIMIT 5;

‎internal/endtoend/testdata/pg_vector/postgresql/pgx/schema.sql

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
CREATE EXTENSION IF NOT EXISTS "vector";
2+
3+
CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));

‎internal/endtoend/testdata/pg_vector/postgresql/pgx/sqlc.yaml

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
version: "2"
2+
sql:
3+
- engine: "postgresql"
4+
schema: "schema.sql"
5+
queries: "query.sql"
6+
gen:
7+
go:
8+
package: "querytest"
9+
out: "go"
10+
sql_package: "pgx/v5"

0 commit comments

Comments
 (0)