From 93669a7fa37dad1651d6adb797eb2fb62a9f785e Mon Sep 17 00:00:00 2001 From: Jille Timmermans Date: Fri, 21 Jan 2022 21:15:47 +0100 Subject: [PATCH] fix(copyfrom): Fix imports when non-copyfrom queries needed imports that copyfrom queries didn't The attached test case is broken without the fix and passes with the fix. --- internal/codegen/golang/imports.go | 27 +++++++++++- .../postgresql/pgx/go/copyfrom.go | 41 +++++++++++++++++++ .../copyfrom_imports/postgresql/pgx/go/db.go | 31 ++++++++++++++ .../postgresql/pgx/go/models.go | 12 ++++++ .../postgresql/pgx/go/query.sql.go | 23 +++++++++++ .../copyfrom_imports/postgresql/pgx/query.sql | 8 ++++ .../copyfrom_imports/postgresql/pgx/sqlc.json | 13 ++++++ 7 files changed, 154 insertions(+), 1 deletion(-) create mode 100644 internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/copyfrom.go create mode 100644 internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/db.go create mode 100644 internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/models.go create mode 100644 internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/query.sql.go create mode 100644 internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/query.sql create mode 100644 internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/sqlc.json diff --git a/internal/codegen/golang/imports.go b/internal/codegen/golang/imports.go index 516cbb3018..93ff4d4376 100644 --- a/internal/codegen/golang/imports.go +++ b/internal/codegen/golang/imports.go @@ -99,7 +99,7 @@ func (i *importer) Imports(filename string) [][]ImportSpec { case querierFileName: return mergeImports(i.interfaceImports()) case copyfromFileName: - return mergeImports(i.interfaceImports()) + return mergeImports(i.copyfromImports()) default: return mergeImports(i.queryImports(filename)) } @@ -367,3 +367,28 @@ func (i *importer) queryImports(filename string) fileImports { return sortedImports(std, pkg) } + +func (i *importer) copyfromImports() fileImports { + std, pkg := buildImports(i.Settings, i.Queries, func(name string) bool { + for _, q := range i.Queries { + if q.Cmd != metadata.CmdCopyFrom { + continue + } + if q.hasRetType() { + if strings.HasPrefix(q.Ret.Type(), name) { + return true + } + } + if !q.Arg.isEmpty() { + if strings.HasPrefix(q.Arg.Type(), name) { + return true + } + } + } + return false + }) + + std["context"] = struct{}{} + + return sortedImports(std, pkg) +} diff --git a/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/copyfrom.go b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/copyfrom.go new file mode 100644 index 0000000000..009ef84092 --- /dev/null +++ b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/copyfrom.go @@ -0,0 +1,41 @@ +// Code generated by sqlc. DO NOT EDIT. +// source: copyfrom.go + +package querytest + +import ( + "context" +) + +// iteratorForInsertValues implements pgx.CopyFromSource. +type iteratorForInsertValues struct { + rows []InsertValuesParams + skippedFirstNextCall bool +} + +func (r *iteratorForInsertValues) Next() bool { + if len(r.rows) == 0 { + return false + } + if !r.skippedFirstNextCall { + r.skippedFirstNextCall = true + return true + } + r.rows = r.rows[1:] + return len(r.rows) > 0 +} + +func (r iteratorForInsertValues) Values() ([]interface{}, error) { + return []interface{}{ + r.rows[0].A, + r.rows[0].B, + }, nil +} + +func (r iteratorForInsertValues) Err() error { + return nil +} + +func (q *Queries) InsertValues(ctx context.Context, arg []InsertValuesParams) (int64, error) { + return q.db.CopyFrom(ctx, []string{"myschema", "foo"}, []string{"a", "b"}, &iteratorForInsertValues{rows: arg}) +} diff --git a/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/db.go b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/db.go new file mode 100644 index 0000000000..062b72ba02 --- /dev/null +++ b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/db.go @@ -0,0 +1,31 @@ +// Code generated by sqlc. DO NOT EDIT. + +package querytest + +import ( + "context" + + "github.com/jackc/pgconn" + "github.com/jackc/pgx/v4" +) + +type DBTX interface { + Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error) + Query(context.Context, string, ...interface{}) (pgx.Rows, error) + QueryRow(context.Context, string, ...interface{}) pgx.Row + CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) +} + +func New(db DBTX) *Queries { + return &Queries{db: db} +} + +type Queries struct { + db DBTX +} + +func (q *Queries) WithTx(tx pgx.Tx) *Queries { + return &Queries{ + db: tx, + } +} diff --git a/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/models.go b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/models.go new file mode 100644 index 0000000000..319a4b26a2 --- /dev/null +++ b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/models.go @@ -0,0 +1,12 @@ +// Code generated by sqlc. DO NOT EDIT. + +package querytest + +import ( + "database/sql" +) + +type MyschemaFoo struct { + A sql.NullString + B sql.NullInt32 +} diff --git a/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/query.sql.go b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/query.sql.go new file mode 100644 index 0000000000..788bd34c7f --- /dev/null +++ b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/go/query.sql.go @@ -0,0 +1,23 @@ +// Code generated by sqlc. DO NOT EDIT. +// source: query.sql + +package querytest + +import ( + "context" + "database/sql" +) + +const insertSingleValue = `-- name: InsertSingleValue :exec +INSERT INTO myschema.foo (a) VALUES ($1) +` + +func (q *Queries) InsertSingleValue(ctx context.Context, a sql.NullString) error { + _, err := q.db.Exec(ctx, insertSingleValue, a) + return err +} + +type InsertValuesParams struct { + A sql.NullString + B sql.NullInt32 +} diff --git a/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/query.sql b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/query.sql new file mode 100644 index 0000000000..86bab7db5a --- /dev/null +++ b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/query.sql @@ -0,0 +1,8 @@ +CREATE SCHEMA myschema; +CREATE TABLE myschema.foo (a text, b integer); + +-- name: InsertValues :copyfrom +INSERT INTO myschema.foo (a, b) VALUES ($1, $2); + +-- name: InsertSingleValue :exec +INSERT INTO myschema.foo (a) VALUES ($1); diff --git a/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/sqlc.json b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/sqlc.json new file mode 100644 index 0000000000..9403bd0279 --- /dev/null +++ b/internal/endtoend/testdata/copyfrom_imports/postgresql/pgx/sqlc.json @@ -0,0 +1,13 @@ +{ + "version": "1", + "packages": [ + { + "path": "go", + "engine": "postgresql", + "sql_package": "pgx/v4", + "name": "querytest", + "schema": "query.sql", + "queries": "query.sql" + } + ] +}