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 56df490e44..59d51e44ca 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,7 +15,9 @@ 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) + 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 e8d53d2b9e..5cd7a85335 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; @@ -38,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 73594f0301..d5a18fa3f4 100644 --- a/examples/batch/postgresql/query.sql.go +++ b/examples/batch/postgresql/query.sql.go @@ -7,28 +7,39 @@ package batch import ( "context" + + "github.com/jackc/pgconn" ) 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 } +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 +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 f0e53ef6f0..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 } @@ -407,11 +410,14 @@ func (i *importer) copyfromImports() fileImports { } func (i *importer) batchImports(filename string) fileImports { - std, pkg := buildImports(i.Settings, i.Queries, func(name string) bool { - for _, q := range i.Queries { - if !usesBatch([]Query{q}) { - continue - } + batchQueries := make([]Query, 0, len(i.Queries)) + for _, q := range i.Queries { + if usesBatch([]Query{q}) { + batchQueries = append(batchQueries, q) + } + } + std, pkg := buildImports(i.Settings, batchQueries, func(name string) bool { + for _, q := range batchQueries { if q.hasRetType() { if q.Ret.EmitStruct() { for _, f := range q.Ret.Struct.Fields { 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() }