From 548aaceffc0b91d909522d162ac9ca61a6a28fee Mon Sep 17 00:00:00 2001 From: xobotyi Date: Fri, 24 Jan 2025 17:37:54 +0100 Subject: [PATCH 1/4] upgrade go version to 1.23 in order to support iterators --- go.mod | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From 812c9373f0cda5fc45fa399d30d459f4a3837d8c Mon Sep 17 00:00:00 2001 From: xobotyi Date: Fri, 24 Jan 2025 18:07:46 +0100 Subject: [PATCH 2/4] implement scanning iterator `AllRowsScanned` --- rows.go | 25 +++++++++++++++ rows_test.go | 89 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/rows.go b/rows.go index 268559ce6..ab2d06455 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 iterator that read and scans rows one-by-one. It closes +// the rows automatically on return. +// +// In case rows.Err() returns non-nil error after all rows are read, it will +// trigger extra yield with 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)) { + break + } + } + + // 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) + } + }) +} From c89724d9e2237c24c0feed65005fea70e083533b Mon Sep 17 00:00:00 2001 From: xobotyi Date: Fri, 24 Jan 2025 20:29:58 +0100 Subject: [PATCH 3/4] `AllRowsScanned` return instead of break Previous approach could cause panic in case loop broken by developer and error occurred on rows. --- rows.go | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/rows.go b/rows.go index ab2d06455..dbf50637d 100644 --- a/rows.go +++ b/rows.go @@ -667,18 +667,21 @@ func RowToAddrOfStructByNameLax[T any](row CollectableRow) (*T, error) { return &value, err } -// AllRowsScanned returns iterator that read and scans rows one-by-one. It closes -// the rows automatically on return. +// AllRowsScanned returns an iterator that reads and scans rows one-by-one. It automatically +// closes the rows when done. // -// In case rows.Err() returns non-nil error after all rows are read, it will -// trigger extra yield with zero value and the error. +// 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. +// +// If the caller's logic implies the possibility of an early loop break, rows.Err() should +// be checked after the loop. 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)) { - break + return } } From 40ff024c0ab791ca1c0ad6cfc08f6e6401c12930 Mon Sep 17 00:00:00 2001 From: xobotyi Date: Sat, 25 Jan 2025 15:11:58 +0100 Subject: [PATCH 4/4] fix `AllRowsScanned` docs regarding error check --- rows.go | 3 --- 1 file changed, 3 deletions(-) diff --git a/rows.go b/rows.go index dbf50637d..81adef3bc 100644 --- a/rows.go +++ b/rows.go @@ -672,9 +672,6 @@ func RowToAddrOfStructByNameLax[T any](row CollectableRow) (*T, error) { // // 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. -// -// If the caller's logic implies the possibility of an early loop break, rows.Err() should -// be checked after the loop. func AllRowsScanned[T any](rows Rows, fn RowToFunc[T]) iter.Seq2[T, error] { return func(yield func(T, error) bool) { defer rows.Close()