From 289098b8b2962317c9fb6b1a41317600b0356a9f Mon Sep 17 00:00:00 2001 From: Sjors Gielen Date: Fri, 16 Oct 2026 17:41:12 +0200 Subject: [PATCH 1/2] Add the :iter query annotation A :iter query generates a method returning an iter.Seq over the query's rows, plus a function returning the query's error once the loop is done. The PostgreSQL authors example shows its use. --- examples/authors/postgresql/db_test.go | 70 +++++++++++++++++++ examples/authors/postgresql/query.sql | 4 ++ examples/authors/postgresql/query.sql.go | 42 +++++++++++ internal/codegen/golang/gen.go | 4 +- internal/codegen/golang/imports.go | 10 +++ internal/codegen/golang/query.go | 2 +- internal/codegen/golang/result.go | 11 +++ .../golang/templates/pgx/interfaceCode.tmpl | 5 ++ .../golang/templates/pgx/queryCode.tmpl | 36 ++++++++++ .../templates/stdlib/interfaceCode.tmpl | 5 ++ .../golang/templates/stdlib/queryCode.tmpl | 39 +++++++++++ internal/metadata/meta.go | 3 +- internal/sql/validate/cmd.go | 2 +- 13 files changed, 228 insertions(+), 5 deletions(-) diff --git a/examples/authors/postgresql/db_test.go b/examples/authors/postgresql/db_test.go index 0fdd8824be..c2403c8a6d 100644 --- a/examples/authors/postgresql/db_test.go +++ b/examples/authors/postgresql/db_test.go @@ -4,6 +4,8 @@ package authors import ( "context" + "errors" + "slices" "testing" "github.com/jackc/pgx/v5" @@ -47,3 +49,71 @@ func TestAuthors(t *testing.T) { } t.Log(fetchedAuthor) } + +func TestIterAuthors(t *testing.T) { + ctx := context.Background() + uri := local.PostgreSQL(t, []string{"schema.sql"}) + db, err := pgx.Connect(ctx, uri) + if err != nil { + t.Fatal(err) + } + defer db.Close(ctx) + + q := New(db) + + var last Author + for _, name := range []string{"Rob Pike", "Brian Kernighan", "Ken Thompson"} { + if last, err = q.CreateAuthor(ctx, CreateAuthorParams{Name: name}); err != nil { + t.Fatal(err) + } + } + + // the query only runs once the iterator is ranged over, so the + // connection is still free for other queries until then + authors, ferr := q.IterAuthors(ctx) + if _, err := q.GetAuthor(ctx, last.ID); err != nil { + t.Fatalf("query before ranging: %v", err) + } + var names []string + for author := range authors { + names = append(names, author.Name) + } + if err := ferr(); err != nil { + t.Fatal(err) + } + if want := []string{"Brian Kernighan", "Ken Thompson", "Rob Pike"}; !slices.Equal(names, want) { + t.Fatalf("got %v, want %v", names, want) + } + + // the iterator is single-use: ranging over it again yields nothing and + // reports an error + for author := range authors { + t.Errorf("second range yielded %v", author) + } + if err := ferr(); err == nil { + t.Error("ranging twice reported no error") + } + + // breaking out of the loop early releases the connection + authors, ferr = q.IterAuthors(ctx) + for range authors { + break + } + if err := ferr(); err != nil { + t.Fatal(err) + } + if _, err := q.GetAuthor(ctx, last.ID); err != nil { + t.Errorf("query after breaking out: %v", err) + } + + // a query that fails reports its error through the returned func + cctx, cancel := context.WithCancel(ctx) + cancel() + authors, ferr = q.IterAuthors(cctx) + for author := range authors { + t.Errorf("cancelled query yielded %v", author) + } + if err := ferr(); !errors.Is(err, context.Canceled) { + t.Errorf("got error %v, want context.Canceled", err) + } +} diff --git a/examples/authors/postgresql/query.sql b/examples/authors/postgresql/query.sql index 75e38b2caf..142019de96 100644 --- a/examples/authors/postgresql/query.sql +++ b/examples/authors/postgresql/query.sql @@ -17,3 +17,7 @@ RETURNING *; -- name: DeleteAuthor :exec DELETE FROM authors WHERE id = $1; + +-- name: IterAuthors :iter +SELECT * FROM authors +ORDER BY name; diff --git a/examples/authors/postgresql/query.sql.go b/examples/authors/postgresql/query.sql.go index 2910a809f3..493606dad7 100644 --- a/examples/authors/postgresql/query.sql.go +++ b/examples/authors/postgresql/query.sql.go @@ -7,6 +7,8 @@ package authors import ( "context" + "errors" + "iter" "github.com/jackc/pgx/v5/pgtype" ) @@ -54,6 +56,46 @@ func (q *Queries) GetAuthor(ctx context.Context, id int64) (Author, error) { return i, err } +const iterAuthors = `-- name: IterAuthors :iter +SELECT id, name, bio FROM authors +ORDER BY name +` + +// IterAuthors returns a single-use iterator. The query runs when the +// iterator is ranged over, and the returned func reports its error +// after the loop. +func (q *Queries) IterAuthors(ctx context.Context) (iter.Seq[Author], func() error) { + var _err error + var _used bool + return func(_yield func(Author) bool) { + if _used { + _err = errors.New("query IterAuthors: iterator already used") + return + } + _used = true + rows, err := q.db.Query(ctx, iterAuthors) + if err != nil { + _err = err + return + } + defer rows.Close() + for rows.Next() { + var i Author + if err := rows.Scan(&i.ID, &i.Name, &i.Bio); err != nil { + _err = err + return + } + if !_yield(i) { + break + } + } + rows.Close() + if err := rows.Err(); err != nil { + _err = err + } + }, func() error { return _err } +} + const listAuthors = `-- name: ListAuthors :many SELECT id, name, bio FROM authors ORDER BY name diff --git a/internal/codegen/golang/gen.go b/internal/codegen/golang/gen.go index 725f5970cf..2f0e276742 100644 --- a/internal/codegen/golang/gen.go +++ b/internal/codegen/golang/gen.go @@ -76,7 +76,7 @@ func (t *tmplCtx) codegenQueryMethod(q Query) string { } return db + ".QueryRowContext" - case ":many": + case ":many", ":iter": if t.EmitPreparedQueries { return "q.query" } @@ -94,7 +94,7 @@ func (t *tmplCtx) codegenQueryRetval(q Query) (string, error) { switch q.Cmd { case ":one": return "row :=", nil - case ":many": + case ":many", ":iter": return "rows, err :=", nil case ":exec": return "_, err :=", nil diff --git a/internal/codegen/golang/imports.go b/internal/codegen/golang/imports.go index e1370e48eb..d63275d49a 100644 --- a/internal/codegen/golang/imports.go +++ b/internal/codegen/golang/imports.go @@ -206,6 +206,9 @@ func buildImports(options *opts.Options, engine string, queries []Query, uses fu std["database/sql"] = struct{}{} } } + if q.Cmd == metadata.CmdIter { + std["iter"] = struct{}{} + } } for typeName, pkg := range stdlibTypes { @@ -446,6 +449,13 @@ func (i *importer) queryImports(filename string) fileImports { } } + // An :iter method reports a second range over its iterator as an error. + for _, q := range gq { + if q.Cmd == metadata.CmdIter { + std["errors"] = struct{}{} + } + } + if anyNonCopyFrom { std["context"] = struct{}{} } diff --git a/internal/codegen/golang/query.go b/internal/codegen/golang/query.go index 6bd9b70e25..b1ffeef709 100644 --- a/internal/codegen/golang/query.go +++ b/internal/codegen/golang/query.go @@ -304,7 +304,7 @@ 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.CmdIter || q.Cmd == metadata.CmdBatchMany || q.Cmd == metadata.CmdBatchOne return scanned && !q.Ret.isEmpty() } diff --git a/internal/codegen/golang/result.go b/internal/codegen/golang/result.go index 82059c54b7..36c5dfb015 100644 --- a/internal/codegen/golang/result.go +++ b/internal/codegen/golang/result.go @@ -205,6 +205,16 @@ func buildQueries(req *plugin.GenerateRequest, options *opts.Options, enums []En } comments := query.Comments + if query.Cmd == metadata.CmdIter { + if len(comments) > 0 { + comments = append(comments, "") + } + comments = append(comments, + " "+query.Name+" returns a single-use iterator. The query runs when the", + " iterator is ranged over, and the returned func reports its error", + " after the loop.", + ) + } if options.EmitSqlAsComment { if len(comments) == 0 { comments = append(comments, query.Name) @@ -370,6 +380,7 @@ func buildQueries(req *plugin.GenerateRequest, options *opts.Options, enums []En var cmdReturnsData = map[string]struct{}{ metadata.CmdBatchMany: {}, metadata.CmdBatchOne: {}, + metadata.CmdIter: {}, metadata.CmdMany: {}, metadata.CmdOne: {}, } diff --git a/internal/codegen/golang/templates/pgx/interfaceCode.tmpl b/internal/codegen/golang/templates/pgx/interfaceCode.tmpl index cf7cd36cb9..cd57a63418 100644 --- a/internal/codegen/golang/templates/pgx/interfaceCode.tmpl +++ b/internal/codegen/golang/templates/pgx/interfaceCode.tmpl @@ -20,6 +20,11 @@ {{end -}} {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error) {{- end}} + {{- if eq .Cmd ":iter" }} + {{range .Comments}}//{{.}} + {{end -}} + {{.MethodName}}(ctx context.Context, {{if $dbtxParam}}db DBTX, {{end}}{{.Arg.Pair}}) (iter.Seq[{{.Ret.DefineType}}], func() error) + {{- end}} {{- if and (eq .Cmd ":exec") ($dbtxParam) }} {{range .Comments}}//{{.}} {{end -}} diff --git a/internal/codegen/golang/templates/pgx/queryCode.tmpl b/internal/codegen/golang/templates/pgx/queryCode.tmpl index 59a88c880a..1dbdce8062 100644 --- a/internal/codegen/golang/templates/pgx/queryCode.tmpl +++ b/internal/codegen/golang/templates/pgx/queryCode.tmpl @@ -79,6 +79,42 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ([]{{.Ret. } {{end}} +{{if eq .Cmd ":iter"}} +{{range .Comments}}//{{.}} +{{end -}} +func (q *Queries) {{.MethodName}}(ctx context.Context, {{if $.EmitMethodsWithDBArgument}}db DBTX, {{end}}{{.Arg.Pair}}) (iter.Seq[{{.Ret.DefineType}}], func() error) { + var _err error + var _used bool + return func(_yield func({{.Ret.DefineType}}) bool) { + if _used { + _err = errors.New("query {{.MethodName}}: iterator already used") + return + } + _used = true + rows, err := {{if $.EmitMethodsWithDBArgument}}db{{else}}q.db{{end}}.Query(ctx, {{.ConstantName}}, {{.Arg.Params}}) + if err != nil { + _err = {{if $.WrapErrors}}fmt.Errorf("query {{.MethodName}}: %w", err){{else}}err{{end}} + return + } + defer rows.Close() + for rows.Next() { + var {{.Ret.Name}} {{.Ret.Type}} + if err := rows.Scan({{.Ret.Scan}}); err != nil { + _err = {{if $.WrapErrors}}fmt.Errorf("query {{.MethodName}}: %w", err){{else}}err{{end}} + return + } + if !_yield({{.Ret.ReturnName}}) { + break + } + } + rows.Close() + if err := rows.Err(); err != nil { + _err = {{if $.WrapErrors}}fmt.Errorf("query {{.MethodName}}: %w", err){{else}}err{{end}} + } + }, func() error { return _err } +} +{{end}} + {{if eq .Cmd ":exec"}} {{range .Comments}}//{{.}} {{end -}} diff --git a/internal/codegen/golang/templates/stdlib/interfaceCode.tmpl b/internal/codegen/golang/templates/stdlib/interfaceCode.tmpl index 3cbefe6df4..52d49d6112 100644 --- a/internal/codegen/golang/templates/stdlib/interfaceCode.tmpl +++ b/internal/codegen/golang/templates/stdlib/interfaceCode.tmpl @@ -20,6 +20,11 @@ {{end -}} {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error) {{- end}} + {{- if eq .Cmd ":iter" }} + {{range .Comments}}//{{.}} + {{end -}} + {{.MethodName}}(ctx context.Context, {{if $dbtxParam}}db DBTX, {{end}}{{.Arg.Pair}}) (iter.Seq[{{.Ret.DefineType}}], func() error) + {{- end}} {{- if and (eq .Cmd ":exec") ($dbtxParam) }} {{range .Comments}}//{{.}} {{end -}} diff --git a/internal/codegen/golang/templates/stdlib/queryCode.tmpl b/internal/codegen/golang/templates/stdlib/queryCode.tmpl index 8fd67f1567..7c4e156166 100644 --- a/internal/codegen/golang/templates/stdlib/queryCode.tmpl +++ b/internal/codegen/golang/templates/stdlib/queryCode.tmpl @@ -68,6 +68,45 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{ dbarg }} {{.Arg.Pair}} } {{end}} +{{if eq .Cmd ":iter"}} +{{range .Comments}}//{{.}} +{{end -}} +func (q *Queries) {{.MethodName}}(ctx context.Context, {{ dbarg }} {{.Arg.Pair}}) (iter.Seq[{{.Ret.DefineType}}], func() error) { + var _err error + var _used bool + return func(_yield func({{.Ret.DefineType}}) bool) { + if _used { + _err = errors.New("query {{.MethodName}}: iterator already used") + return + } + _used = true + {{ template "queryCodeStdExec" . }} + if err != nil { + _err = {{if $.WrapErrors}}fmt.Errorf("query {{.MethodName}}: %w", err){{else}}err{{end}} + return + } + defer rows.Close() + for rows.Next() { + var {{.Ret.Name}} {{.Ret.Type}} + if err := rows.Scan({{.Ret.Scan}}); err != nil { + _err = {{if $.WrapErrors}}fmt.Errorf("query {{.MethodName}}: %w", err){{else}}err{{end}} + return + } + if !_yield({{.Ret.ReturnName}}) { + break + } + } + if err := rows.Close(); err != nil { + _err = {{if $.WrapErrors}}fmt.Errorf("query {{.MethodName}}: %w", err){{else}}err{{end}} + return + } + if err := rows.Err(); err != nil { + _err = {{if $.WrapErrors}}fmt.Errorf("query {{.MethodName}}: %w", err){{else}}err{{end}} + } + }, func() error { return _err } +} +{{end}} + {{if eq .Cmd ":exec"}} {{range .Comments}}//{{.}} {{end -}} diff --git a/internal/metadata/meta.go b/internal/metadata/meta.go index 76ee992a7a..ab2baa5223 100644 --- a/internal/metadata/meta.go +++ b/internal/metadata/meta.go @@ -32,6 +32,7 @@ const ( CmdExecRows = ":execrows" CmdExecLastId = ":execlastid" CmdMany = ":many" + CmdIter = ":iter" CmdOne = ":one" CmdCopyFrom = ":copyfrom" CmdBatchExec = ":batchexec" @@ -106,7 +107,7 @@ func ParseQueryNameAndType(t string, commentStyle CommentSyntax) (string, string queryName := part[2] queryType := strings.TrimSpace(part[3]) switch queryType { - case CmdOne, CmdMany, CmdExec, CmdExecResult, CmdExecRows, CmdExecLastId, CmdCopyFrom, CmdBatchExec, CmdBatchMany, CmdBatchOne: + case CmdOne, CmdMany, CmdIter, CmdExec, CmdExecResult, CmdExecRows, CmdExecLastId, CmdCopyFrom, CmdBatchExec, CmdBatchMany, CmdBatchOne: default: return "", "", fmt.Errorf("invalid query type: %s", queryType) } diff --git a/internal/sql/validate/cmd.go b/internal/sql/validate/cmd.go index e9d9dd03cb..702379683c 100644 --- a/internal/sql/validate/cmd.go +++ b/internal/sql/validate/cmd.go @@ -65,7 +65,7 @@ func Cmd(n ast.Node, name, cmd string) error { return err } } - if !(cmd == metadata.CmdMany || cmd == metadata.CmdOne || cmd == metadata.CmdBatchMany || cmd == metadata.CmdBatchOne) { + if !(cmd == metadata.CmdMany || cmd == metadata.CmdIter || cmd == metadata.CmdOne || cmd == metadata.CmdBatchMany || cmd == metadata.CmdBatchOne) { return nil } var list *ast.List From 138b1c9f25d5eff50473326480a8faff50def14b Mon Sep 17 00:00:00 2001 From: Sjors Gielen Date: Fri, 16 Oct 2026 19:15:21 +0200 Subject: [PATCH 2/2] Add documentation for the :iter query annotation --- docs/reference/query-annotations.md | 43 +++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/docs/reference/query-annotations.md b/docs/reference/query-annotations.md index 4fabe05aae..4abdbb58c6 100644 --- a/docs/reference/query-annotations.md +++ b/docs/reference/query-annotations.md @@ -77,6 +77,49 @@ func (q *Queries) InsertAuthor(ctx context.Context, name string) (int64, error) } ``` +## `:iter` + +The generated method will return an [iter.Seq](https://pkg.go.dev/iter#Seq) that +streams records as they are scanned, together with a function that eventually +reports the query's error. This needs Go 1.23 or later. + +```sql +-- name: ListAuthors :iter +SELECT * FROM authors +ORDER BY name; +``` + +```go +func (q *Queries) ListAuthors(ctx context.Context) (iter.Seq[Author], func() error) { + // ... +} +``` + +To consume, you can range over the `iter.Seq`: + +```go +authors, ferr := q.ListAuthors(ctx) +for author := range authors { + // ... +} +if err := ferr(); err != nil { + // ... +} +``` + +The query starts only when the iterator is ranged over, not when the method is +called. When `q` runs on a single connection or transaction, it must not be +used inside the `for` loop, as the query is still reading from it. Once the +loop ends, by running out of rows or by `break`, the rows are closed and the +connection is released. + +The function returned alongside the iterator reports any error from running the +query or scanning its rows, so call it after the loop. + +The iterator is single-use: it runs its query once, which matters for a +`DELETE ... RETURNING` or `INSERT ... RETURNING`. Ranging over it a second +time yields nothing, and the returned function then reports an error. + ## `:many` The generated method will return a slice of records via