diff --git a/go.mod b/go.mod index 0378b800c..2c02f58fa 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/jackc/pgx/v5 -go 1.21 +go 1.23 require ( github.com/jackc/pgpassfile v1.0.0 diff --git a/rows.go b/rows.go index 268559ce6..81adef3bc 100644 --- a/rows.go +++ b/rows.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "iter" "reflect" "strings" "sync" @@ -666,6 +667,30 @@ func RowToAddrOfStructByNameLax[T any](row CollectableRow) (*T, error) { return &value, err } +// AllRowsScanned returns an iterator that reads and scans rows one-by-one. It automatically +// closes the rows when done. +// +// If rows.Err() returns a non-nil error after all rows are read, it will trigger an extra +// yield with a zero value and the error. +func AllRowsScanned[T any](rows Rows, fn RowToFunc[T]) iter.Seq2[T, error] { + return func(yield func(T, error) bool) { + defer rows.Close() + + for rows.Next() { + if !yield(fn(rows)) { + return + } + } + + // we don't have another choice but to push one more time + // in order to propagate the error to user + if err := rows.Err(); err != nil { + var zero T + yield(zero, err) + } + } +} + type namedStructRowScanner struct { ptrToStruct any lax bool diff --git a/rows_test.go b/rows_test.go index 4cda957fc..a789f5a0f 100644 --- a/rows_test.go +++ b/rows_test.go @@ -993,3 +993,92 @@ insert into products (name, price) values // Fries: $5 // Soft Drink: $3 } + +func ExampleAllRowsScanned() { + ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) + defer cancel() + + conn, err := pgx.Connect(ctx, os.Getenv("PGX_TEST_DATABASE")) + if err != nil { + fmt.Printf("Unable to establish connection: %v", err) + return + } + + if conn.PgConn().ParameterStatus("crdb_version") != "" { + // Skip test / example when running on CockroachDB. Since an example can't be skipped fake success instead. + fmt.Println(`Cheeseburger: $10 +Fries: $5 +Soft Drink: $3`) + return + } + + // Setup example schema and data. + _, err = conn.Exec(ctx, ` +create temporary table products ( + id int primary key generated by default as identity, + name varchar(100) not null, + price int not null +); + +insert into products (name, price) values + ('Cheeseburger', 10), + ('Double Cheeseburger', 14), + ('Fries', 5), + ('Soft Drink', 3); +`) + if err != nil { + fmt.Printf("Unable to setup example schema and data: %v", err) + return + } + + type product struct { + ID int32 + Name string + Type string + Price int32 + } + + result := make([]product, 0, 3) + + rows, _ := conn.Query(ctx, "select * from products where price < $1 order by price desc", 12) + for row, err := range pgx.AllRowsScanned[product](rows, pgx.RowToStructByNameLax) { + if err != nil { + fmt.Printf("AllRowsScanned error: %v", err) + return + } + + // our business logic here + result = append(result, row) + } + + for _, p := range result { + fmt.Printf("%s: $%d\n", p.Name, p.Price) + } + + // Output: + // Cheeseburger: $10 + // Fries: $5 + // Soft Drink: $3 +} + +func TestAllRowsScanned(t *testing.T) { + defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) { + type resultRow struct { + N int32 `db:"n"` + } + + rows, _ := conn.Query(ctx, `select n from generate_series(0, 99) n`) + + results := make([]resultRow, 0, 100) + + for row, err := range pgx.AllRowsScanned[resultRow](rows, pgx.RowToStructByName) { + require.NoError(t, err) + results = append(results, row) + } + + assert.Len(t, results, 100) + for i := range results { + assert.Equal(t, int32(i), results[i].N) + } + }) +}