Add tests for PR #4646 (:iter query annotation)

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

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)
+	}
+}

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;

examples/authors/sqlite/db_test.go

@@ -5,6 +5,8 @@ package authors
 import (
 	"context"
 	"database/sql"
+	"errors"
+	"slices"
 	"testing"
 
 	"github.com/sqlc-dev/sqlc/internal/sqltest"
@@ -46,3 +48,66 @@ func TestAuthors(t *testing.T) {
 	}
 	t.Log(fetchedAuthor)
 }
+
+func TestIterAuthors(t *testing.T) {
+	sdb, cleanup := sqltest.SQLite(t, []string{"schema.sql"})
+	defer sdb.Close()
+	defer cleanup()
+
+	ctx := context.Background()
+	db := New(sdb)
+
+	for _, name := range []string{"Rob Pike", "Brian Kernighan", "Ken Thompson"} {
+		if _, err := db.CreateAuthor(ctx, CreateAuthorParams{Name: name}); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	// the query only runs once the iterator is ranged over
+	authors, ferr := db.IterAuthors(ctx)
+	if inUse := sdb.Stats().InUse; inUse != 0 {
+		t.Fatalf("%d connections in use before ranging, want 0", inUse)
+	}
+	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 = db.IterAuthors(ctx)
+	for range authors {
+		break
+	}
+	if err := ferr(); err != nil {
+		t.Fatal(err)
+	}
+	if inUse := sdb.Stats().InUse; inUse != 0 {
+		t.Errorf("%d connections in use after breaking out, want 0", inUse)
+	}
+
+	// a query that fails reports its error through the returned func
+	cctx, cancel := context.WithCancel(ctx)
+	cancel()
+	authors, ferr = db.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)
+	}
+}

examples/authors/sqlite/query.sql

@@ -16,3 +16,7 @@ INSERT INTO authors (
 /* name: DeleteAuthor :exec */
 DELETE FROM authors
 WHERE id = ?;
+
+/* name: IterAuthors :iter */
+SELECT * FROM authors
+ORDER BY name;

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

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{}{}
 	}

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()
 }