From 6d9ccfb9d03b55141b3c583252ac4fbf01663085 Mon Sep 17 00:00:00 2001 From: Chris Lee Date: Thu, 8 Sep 2022 23:48:55 -0700 Subject: [PATCH 1/3] fix: use only batch queries to build batch imports --- examples/batch/postgresql/querier.go | 3 +++ examples/batch/postgresql/query.sql | 4 ++++ examples/batch/postgresql/query.sql.go | 11 +++++++++++ internal/codegen/golang/imports.go | 8 +++++++- 4 files changed, 25 insertions(+), 1 deletion(-) diff --git a/examples/batch/postgresql/querier.go b/examples/batch/postgresql/querier.go index 56df490e44..86ecae489d 100644 --- a/examples/batch/postgresql/querier.go +++ b/examples/batch/postgresql/querier.go @@ -6,6 +6,8 @@ package batch import ( "context" + + "github.com/jackc/pgconn" ) type Querier interface { @@ -13,6 +15,7 @@ type Querier interface { CreateAuthor(ctx context.Context, name string) (Author, error) CreateBook(ctx context.Context, arg []CreateBookParams) *CreateBookBatchResults DeleteBook(ctx context.Context, bookID []int32) *DeleteBookBatchResults + DeleteBookExecResult(ctx context.Context, bookID int32) (pgconn.CommandTag, error) GetAuthor(ctx context.Context, authorID int32) (Author, error) UpdateBook(ctx context.Context, arg []UpdateBookParams) *UpdateBookBatchResults } diff --git a/examples/batch/postgresql/query.sql b/examples/batch/postgresql/query.sql index e8d53d2b9e..d65e4308e7 100644 --- a/examples/batch/postgresql/query.sql +++ b/examples/batch/postgresql/query.sql @@ -2,6 +2,10 @@ SELECT * FROM authors WHERE author_id = $1; +-- name: DeleteBookExecResult :execresult +DELETE FROM books +WHERE book_id = $1; + -- name: DeleteBook :batchexec DELETE FROM books WHERE book_id = $1; diff --git a/examples/batch/postgresql/query.sql.go b/examples/batch/postgresql/query.sql.go index 73594f0301..bb1eb1c75b 100644 --- a/examples/batch/postgresql/query.sql.go +++ b/examples/batch/postgresql/query.sql.go @@ -7,6 +7,8 @@ package batch import ( "context" + + "github.com/jackc/pgconn" ) const createAuthor = `-- name: CreateAuthor :one @@ -21,6 +23,15 @@ func (q *Queries) CreateAuthor(ctx context.Context, name string) (Author, error) return i, err } +const deleteBookExecResult = `-- name: DeleteBookExecResult :execresult +DELETE FROM books +WHERE book_id = $1 +` + +func (q *Queries) DeleteBookExecResult(ctx context.Context, bookID int32) (pgconn.CommandTag, error) { + return q.db.Exec(ctx, deleteBookExecResult, bookID) +} + const getAuthor = `-- name: GetAuthor :one SELECT author_id, name FROM authors WHERE author_id = $1 diff --git a/internal/codegen/golang/imports.go b/internal/codegen/golang/imports.go index f0e53ef6f0..1f7954063c 100644 --- a/internal/codegen/golang/imports.go +++ b/internal/codegen/golang/imports.go @@ -407,7 +407,13 @@ func (i *importer) copyfromImports() fileImports { } func (i *importer) batchImports(filename string) fileImports { - std, pkg := buildImports(i.Settings, i.Queries, func(name string) bool { + batchQueries := make([]Query, 0, len(i.Queries)) + for _, q := range i.Queries { + if q.Cmd == metadata.CmdBatchExec || q.Cmd == metadata.CmdBatchMany || q.Cmd == metadata.CmdBatchOne { + batchQueries = append(batchQueries, q) + } + } + std, pkg := buildImports(i.Settings, batchQueries, func(name string) bool { for _, q := range i.Queries { if !usesBatch([]Query{q}) { continue From 608c2f4a69e87d9f66759d46649c0487478a11c2 Mon Sep 17 00:00:00 2001 From: Chris Lee Date: Fri, 9 Sep 2022 00:01:51 -0700 Subject: [PATCH 2/3] reuse usesBatch function instead of checking individually --- internal/codegen/golang/imports.go | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/internal/codegen/golang/imports.go b/internal/codegen/golang/imports.go index 1f7954063c..c3aaab6436 100644 --- a/internal/codegen/golang/imports.go +++ b/internal/codegen/golang/imports.go @@ -409,15 +409,12 @@ func (i *importer) copyfromImports() fileImports { func (i *importer) batchImports(filename string) fileImports { batchQueries := make([]Query, 0, len(i.Queries)) for _, q := range i.Queries { - if q.Cmd == metadata.CmdBatchExec || q.Cmd == metadata.CmdBatchMany || q.Cmd == metadata.CmdBatchOne { + if usesBatch([]Query{q}) { batchQueries = append(batchQueries, q) } } std, pkg := buildImports(i.Settings, batchQueries, func(name string) bool { - for _, q := range i.Queries { - if !usesBatch([]Query{q}) { - continue - } + for _, q := range batchQueries { if q.hasRetType() { if q.Ret.EmitStruct() { for _, f := range q.Ret.Struct.Fields { From 981cc56ba1001d77b3a35aebf29d7cf7ea7d51ae Mon Sep 17 00:00:00 2001 From: Chris Lee Date: Fri, 9 Sep 2022 00:32:50 -0700 Subject: [PATCH 3/3] fix: add batch cmd types to hasRetTypes --- examples/batch/postgresql/batch.go | 47 ++++++++++++++++++++++++++ examples/batch/postgresql/models.go | 7 ++-- examples/batch/postgresql/querier.go | 1 + examples/batch/postgresql/query.sql | 4 +++ examples/batch/postgresql/query.sql.go | 8 ++--- examples/batch/postgresql/schema.sql | 3 +- internal/codegen/golang/imports.go | 3 ++ internal/codegen/golang/query.go | 3 +- 8 files changed, 68 insertions(+), 8 deletions(-) diff --git a/examples/batch/postgresql/batch.go b/examples/batch/postgresql/batch.go index 09b36019d0..b2958cbb22 100644 --- a/examples/batch/postgresql/batch.go +++ b/examples/batch/postgresql/batch.go @@ -10,6 +10,7 @@ import ( "errors" "time" + "github.com/jackc/pgtype" "github.com/jackc/pgx/v4" ) @@ -212,6 +213,52 @@ func (b *DeleteBookBatchResults) Close() error { return b.br.Close() } +const getBiography = `-- name: GetBiography :batchone +SELECT biography FROM authors +WHERE author_id = $1 +` + +type GetBiographyBatchResults struct { + br pgx.BatchResults + tot int + closed bool +} + +func (q *Queries) GetBiography(ctx context.Context, authorID []int32) *GetBiographyBatchResults { + batch := &pgx.Batch{} + for _, a := range authorID { + vals := []interface{}{ + a, + } + batch.Queue(getBiography, vals...) + } + br := q.db.SendBatch(ctx, batch) + return &GetBiographyBatchResults{br, len(authorID), false} +} + +func (b *GetBiographyBatchResults) QueryRow(f func(int, pgtype.JSONB, error)) { + defer b.br.Close() + for t := 0; t < b.tot; t++ { + var biography pgtype.JSONB + if b.closed { + if f != nil { + f(t, biography, errors.New("batch already closed")) + } + continue + } + row := b.br.QueryRow() + err := row.Scan(&biography) + if f != nil { + f(t, biography, err) + } + } +} + +func (b *GetBiographyBatchResults) Close() error { + b.closed = true + return b.br.Close() +} + const updateBook = `-- name: UpdateBook :batchexec UPDATE books SET title = $1, tags = $2 diff --git a/examples/batch/postgresql/models.go b/examples/batch/postgresql/models.go index 34e3d66afb..ddb3c1370e 100644 --- a/examples/batch/postgresql/models.go +++ b/examples/batch/postgresql/models.go @@ -8,6 +8,8 @@ import ( "database/sql/driver" "fmt" "time" + + "github.com/jackc/pgtype" ) type BookType string @@ -53,8 +55,9 @@ func (ns NullBookType) Value() (driver.Value, error) { } type Author struct { - AuthorID int32 `json:"author_id"` - Name string `json:"name"` + AuthorID int32 `json:"author_id"` + Name string `json:"name"` + Biography pgtype.JSONB `json:"biography"` } type Book struct { diff --git a/examples/batch/postgresql/querier.go b/examples/batch/postgresql/querier.go index 86ecae489d..59d51e44ca 100644 --- a/examples/batch/postgresql/querier.go +++ b/examples/batch/postgresql/querier.go @@ -17,6 +17,7 @@ type Querier interface { DeleteBook(ctx context.Context, bookID []int32) *DeleteBookBatchResults DeleteBookExecResult(ctx context.Context, bookID int32) (pgconn.CommandTag, error) GetAuthor(ctx context.Context, authorID int32) (Author, error) + GetBiography(ctx context.Context, authorID []int32) *GetBiographyBatchResults UpdateBook(ctx context.Context, arg []UpdateBookParams) *UpdateBookBatchResults } diff --git a/examples/batch/postgresql/query.sql b/examples/batch/postgresql/query.sql index d65e4308e7..5cd7a85335 100644 --- a/examples/batch/postgresql/query.sql +++ b/examples/batch/postgresql/query.sql @@ -42,3 +42,7 @@ RETURNING *; UPDATE books SET title = $1, tags = $2 WHERE book_id = $3; + +-- name: GetBiography :batchone +SELECT biography FROM authors +WHERE author_id = $1; diff --git a/examples/batch/postgresql/query.sql.go b/examples/batch/postgresql/query.sql.go index bb1eb1c75b..d5a18fa3f4 100644 --- a/examples/batch/postgresql/query.sql.go +++ b/examples/batch/postgresql/query.sql.go @@ -13,13 +13,13 @@ import ( const createAuthor = `-- name: CreateAuthor :one INSERT INTO authors (name) VALUES ($1) -RETURNING author_id, name +RETURNING author_id, name, biography ` func (q *Queries) CreateAuthor(ctx context.Context, name string) (Author, error) { row := q.db.QueryRow(ctx, createAuthor, name) var i Author - err := row.Scan(&i.AuthorID, &i.Name) + err := row.Scan(&i.AuthorID, &i.Name, &i.Biography) return i, err } @@ -33,13 +33,13 @@ func (q *Queries) DeleteBookExecResult(ctx context.Context, bookID int32) (pgcon } const getAuthor = `-- name: GetAuthor :one -SELECT author_id, name FROM authors +SELECT author_id, name, biography FROM authors WHERE author_id = $1 ` func (q *Queries) GetAuthor(ctx context.Context, authorID int32) (Author, error) { row := q.db.QueryRow(ctx, getAuthor, authorID) var i Author - err := row.Scan(&i.AuthorID, &i.Name) + err := row.Scan(&i.AuthorID, &i.Name, &i.Biography) return i, err } diff --git a/examples/batch/postgresql/schema.sql b/examples/batch/postgresql/schema.sql index 700b6658e7..b79f773f27 100644 --- a/examples/batch/postgresql/schema.sql +++ b/examples/batch/postgresql/schema.sql @@ -1,6 +1,7 @@ CREATE TABLE authors ( author_id SERIAL PRIMARY KEY, - name text NOT NULL DEFAULT '' + name text NOT NULL DEFAULT '', + biography JSONB ); CREATE TYPE book_type AS ENUM ( diff --git a/internal/codegen/golang/imports.go b/internal/codegen/golang/imports.go index c3aaab6436..9937e38359 100644 --- a/internal/codegen/golang/imports.go +++ b/internal/codegen/golang/imports.go @@ -239,6 +239,9 @@ func (i *importer) interfaceImports() fileImports { std, pkg := buildImports(i.Settings, i.Queries, func(name string) bool { for _, q := range i.Queries { if q.hasRetType() { + if usesBatch([]Query{q}) { + continue + } if strings.HasPrefix(q.Ret.Type(), name) { return true } diff --git a/internal/codegen/golang/query.go b/internal/codegen/golang/query.go index 2eed835c6f..430dfb3b8f 100644 --- a/internal/codegen/golang/query.go +++ b/internal/codegen/golang/query.go @@ -165,7 +165,8 @@ type Query struct { } func (q Query) hasRetType() bool { - scanned := q.Cmd == metadata.CmdOne || q.Cmd == metadata.CmdMany + scanned := q.Cmd == metadata.CmdOne || q.Cmd == metadata.CmdMany || + q.Cmd == metadata.CmdBatchMany || q.Cmd == metadata.CmdBatchOne return scanned && !q.Ret.isEmpty() }