feat(codegen): Add support for pgvector and pgvector-go (#2888)

* feat(codegen): Add support for pgvector-go
* test: Update ddl_tests to use FindTests

internal/codegen/golang/imports.go

@@ -219,6 +219,10 @@ func buildImports(settings *plugin.Settings, options *opts, queries []Query, use
  if uses("uuid.NullUUID") && !overrideNullUUID {
  pkg[ImportSpec{Path: "github.com/google/uuid"}] = struct{}{}
  }
+ _, overrideVector := overrideTypes["pgvector.Vector"]
+ if uses("pgvector.Vector") && !overrideVector {
+ pkg[ImportSpec{Path: "github.com/pgvector/pgvector-go"}] = struct{}{}
+ }
 
  // Custom imports
  for _, o := range settings.Overrides {

internal/codegen/golang/postgresql_type.go

@@ -537,6 +537,15 @@ func postgresType(req *plugin.CodeGenRequest, options *opts, col *plugin.Column)
  return "pgtype.Polygon"
  }
 
+ case "vector":
+ if driver == SQLDriverPGXV5 {
+ if emitPointersForNull {
+ return "*pgvector.Vector"
+ } else {
+ return "pgvector.Vector"
+ }
+ }
+
  case "void":
  // A void value can only be scanned into an empty interface.
  return "interface{}"

internal/endtoend/ddl_test.go

@@ -3,9 +3,9 @@ package main
 import (
  "context"
  "fmt"
- "io/fs"
  "os"
  "path/filepath"
+ "slices"
  "strings"
  "testing"
 
@@ -32,28 +32,20 @@ func TestValidSchema(t *testing.T) {
  t.Fatal(err)
  }
 
- files := []string{}
+ for _, replay := range FindTests(t, "testdata", "managed-db") {
+ replay := replay // https://golang.org/doc/faq#closures_and_goroutines
 
- // Find all tests that do not have a stderr.txt file
- err = filepath.Walk("testdata", func(path string, info fs.FileInfo, err error) error {
- if err != nil {
- return err
+ if len(replay.Stderr) > 0 {
+ continue
  }
- if filepath.Base(path) == "sqlc.json" || filepath.Base(path) == "sqlc.yaml" || filepath.Base(path) == "sqlc.yml" {
-  stderr := filepath.Join(filepath.Dir(path), "stderr.txt")
- if _, err := os.Stat(stderr); !os.IsNotExist(err) {
- return nil
+
+ if replay.Exec != nil {
+ if !slices.Contains(replay.Exec.Contexts, "managed-db") {
+ continue
  }
- files = append(files, path)
  }
- return nil
- })
- if err != nil {
- t.Fatal(err)
- }
 
- for _, file := range files {
- file := file // https://golang.org/doc/faq#closures_and_goroutines
+ file := filepath.Join(replay.Path, replay.ConfigName)
  rd, err := os.Open(file)
  if err != nil {
  t.Fatal(err)

internal/endtoend/testdata/go.mod

@@ -10,7 +10,7 @@ require (
  github.com/jackc/pgconn v1.5.1-0.20200601181101-fa742c524853
  github.com/jackc/pgtype v1.6.2
  github.com/jackc/pgx/v4 v4.6.1-0.20200606145419-4e5062306904
- github.com/jackc/pgx/v5 v5.0.1
+ github.com/jackc/pgx/v5 v5.4.3
  github.com/lib/pq v1.9.0
  github.com/sqlc-dev/pqtype v0.2.0
  github.com/sqlc-dev/sqlc-testdata v1.0.0
@@ -24,11 +24,12 @@ require (
  github.com/jackc/pgio v1.0.0 // indirect
  github.com/jackc/pgpassfile v1.0.0 // indirect
  github.com/jackc/pgproto3/v2 v2.0.1 // indirect
- github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b // indirect
+ github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
+ github.com/pgvector/pgvector-go v0.1.1 // indirect
  github.com/volatiletech/inflect v0.0.1 // indirect
  github.com/volatiletech/randomize v0.0.1 // indirect
  github.com/volatiletech/strmangle v0.0.1 // indirect
- golang.org/x/crypto v0.0.0-20220829220503-c86fa9a7ed90 // indirect
- golang.org/x/text v0.3.7 // indirect
- golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 // indirect
+ golang.org/x/crypto v0.9.0 // indirect
+ golang.org/x/text v0.9.0 // indirect
+ golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
 )

internal/endtoend/testdata/go.sum

@@ -47,6 +47,8 @@ github.com/jackc/pgproto3/v2 v2.0.1/go.mod h1:WfJCnwN3HIg9Ish/j3sgWXnAfK8A9Y0bwX
 github.com/jackc/pgservicefile v0.0.0-20200307190119-3430c5407db8/go.mod h1:vsD4gTJCa9TptPL8sPkXrLZ+hDuNrZCnj29CQpr4X1E=
 github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b h1:C8S2+VttkHFdOOCXJe+YGfa4vHYwlt4Zx+IVXQ97jYg=
 github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b/go.mod h1:vsD4gTJCa9TptPL8sPkXrLZ+hDuNrZCnj29CQpr4X1E=
+github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk=
+github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
 github.com/jackc/pgtype v0.0.0-20190421001408-4ed0de4755e0/go.mod h1:hdSHsc1V01CGwFsrv11mJRHWJ6aifDLfdV3aVjFF0zg=
 github.com/jackc/pgtype v0.0.0-20190824184912-ab885b375b90/go.mod h1:KcahbBH1nCMSo2DXpzsoWOAfFkdEtEJpPbVLq8eE+mc=
 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
 github.com/jackc/pgx/v4 v4.6.1-0.20200606145419-4e5062306904/go.mod h1:ZDaNWkt9sW1JMiNn0kdYBaLelIhw7Pg4qd+Vk6tw7Hg=
 github.com/jackc/pgx/v5 v5.0.1 h1:JZu9othr7l8so2JMDAGeDUMXqERAuZpovyfl4H50tdg=
 github.com/jackc/pgx/v5 v5.0.1/go.mod h1:JBbvW3Hdw77jKl9uJrEDATUZIFM2VFPzRq4RWIhkF4o=
+github.com/jackc/pgx/v5 v5.3.1 h1:Fcr8QJ1ZeLi5zsPZqQeUZhNhxfkkKBOgJuYkJHoBOtU=
+github.com/jackc/pgx/v5 v5.3.1/go.mod h1:t3JDKnCBlYIc0ewLF0Q7B8MXmoIaBOZj/ic7iHozM/8=
+github.com/jackc/pgx/v5 v5.4.3 h1:cxFyXhxlvAifxnkKKdlxv8XqUf59tDlYjnV5YYfsJJY=
+github.com/jackc/pgx/v5 v5.4.3/go.mod h1:Ig06C2Vu0t5qXC60W8sqIthScaEnFvojjj9dSljmHRA=
 github.com/jackc/puddle v0.0.0-20190413234325-e4ced69a3a2b/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
 github.com/jackc/puddle v0.0.0-20190608224051-11cab39313c9/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
 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
 github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
 github.com/mattn/go-isatty v0.0.9/go.mod h1:YNRxwqDuOph6SZLI9vUUz6OYw3QyUt7WiY2yME+cCiQ=
 github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
+github.com/pgvector/pgvector-go v0.1.1 h1:kqJigGctFnlWvskUiYIvJRNwUtQl/aMSUZVs0YWQe+g=
+github.com/pgvector/pgvector-go v0.1.1/go.mod h1:wLJgD/ODkdtd2LJK4l6evHXTuG+8PxymYAVomKHOWac=
 github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 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
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
 github.com/volatiletech/inflect v0.0.1 h1:2a6FcMQyhmPZcLa+uet3VJ8gLn/9svWhJxJYwvE8KsU=
 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
 golang.org/x/crypto v0.0.0-20200323165209-0ec3e9974c59/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20220829220503-c86fa9a7ed90 h1:Y/gsMcFOcR+6S6f3YeMKl5g+dZMEWqcz5Czj/GWYbkM=
 golang.org/x/crypto v0.0.0-20220829220503-c86fa9a7ed90/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
+golang.org/x/crypto v0.6.0/go.mod h1:OFC/31mSvZgRz0V1QTNCzfAI1aIRzbiufJtkMIlEp58=
+golang.org/x/crypto v0.9.0 h1:LF6fAI+IutBocDJ2OT0Q1g8plpYljMZ4+lty+dsqw3g=
+golang.org/x/crypto v0.9.0/go.mod h1:yrmDGqONDYtNj3tH8X9dzUun2m2lzPa9ngI6/RUPGR0=
 golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
 golang.org/x/mod v0.0.0-20190513183733-4bf6d317e70e/go.mod h1:mXi4GBBbnImb6dmsKGUJ2LatrhH/nqhxcFungHvyanc=
 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=
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
 golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
+golang.org/x/text v0.7.0 h1:4BRB4x83lYWy72KwLD/qYDuTu7q9PjSagHvijDw7cLo=
+golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
+golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
+golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 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
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=
+golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/guregu/null.v4 v4.0.0 h1:1Wm3S1WEA2I26Kq+6vcW+w0gcDo44YKYD7YIEJNHDjg=
 gopkg.in/guregu/null.v4 v4.0.0/go.mod h1:YoQhUrADuG3i9WqesrCmpNRwm1ypAgSHYqoOcTu/JrI=
 gopkg.in/inconshreveable/log15.v2 v2.0.0-20180818164646-67afb5ed74ec/go.mod h1:aPpfJ7XW+gOuirDoZ8gHhLh3kZ1B08FtV2bbmy7Jv3s=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=

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

@@ -0,0 +1,3 @@
+{
+ "contexts": ["base"]
+}

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

@@ -0,0 +1,8 @@
+-- name: InsertVector :exec
+INSERT INTO items (embedding) VALUES ($1);
+
+-- name: NearestNeighbor :many
+SELECT *
+FROM items
+ORDER BY embedding <-> $1
+LIMIT 5;

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

@@ -0,0 +1,3 @@
+CREATE EXTENSION IF NOT EXISTS "vector";
+
+CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));

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

@@ -0,0 +1,10 @@
+version: "2"
+sql:
+ - engine: "postgresql"
+ schema: "schema.sql"
+ queries: "query.sql"
+ gen:
+ go:
+ package: "querytest"
+ out: "go"
+ sql_package: "pgx/v5"