Add iterator method for accounts and update tests

Author: klaidliadon

table.go

@@ -2,7 +2,9 @@ package pgkit
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"iter"
 	"slices"
 	"time"
 
@@ -156,20 +158,25 @@ func (t *Table[T, PT, IDT]) saveAll(ctx context.Context, records []PT) error {
 	return nil
 }
 
-// Get returns the first record matching the condition.
-func (t *Table[T, PT, IDT]) Get(ctx context.Context, where sq.Sqlizer, orderBy []string) (PT, error) {
+// getListQuery builds a base select query for listing records.
+func (t *Table[T, PT, IDT]) getListQuery(where sq.Sqlizer, orderBy []string) sq.SelectBuilder {
 	if len(orderBy) == 0 {
 		orderBy = []string{t.IDColumn}
 	}
 
-	record := new(T)
-
 	q := t.SQL.
 		Select("*").
 		From(t.Name).
 		Where(where).
-		Limit(1).
 		OrderBy(orderBy...)
+	return q
+}
+
+// Get returns the first record matching the condition.
+func (t *Table[T, PT, IDT]) Get(ctx context.Context, where sq.Sqlizer, orderBy []string) (PT, error) {
+	record := new(T)
+
+	q := t.getListQuery(where, orderBy).Limit(1)
 
 	if err := t.Query.GetOne(ctx, q, record); err != nil {
 		return nil, fmt.Errorf("get record: %w", err)
@@ -180,16 +187,7 @@ func (t *Table[T, PT, IDT]) Get(ctx context.Context, where sq.Sqlizer, orderBy [
 
 // List returns all records matching the condition.
 func (t *Table[T, PT, IDT]) List(ctx context.Context, where sq.Sqlizer, orderBy []string) ([]PT, error) {
-	if len(orderBy) == 0 {
-		orderBy = []string{t.IDColumn}
-	}
-
-	q := t.SQL.
-		Select("*").
-		From(t.Name).
-		Where(where).
-		OrderBy(orderBy...)
-
+	q := t.getListQuery(where, orderBy)
 	var records []PT
 	if err := t.Query.GetAll(ctx, q, &records); err != nil {
 		return nil, err
@@ -198,6 +196,31 @@ func (t *Table[T, PT, IDT]) List(ctx context.Context, where sq.Sqlizer, orderBy
 	return records, nil
 }
 
+// Iter returns an iterator for records matching the condition.
+func (t *Table[T, PT, IDT]) Iter(ctx context.Context, where sq.Sqlizer, orderBy []string) (iter.Seq2[PT, error], error) {
+	q := t.getListQuery(where, orderBy)
+	rows, err := t.Query.QueryRows(ctx, q)
+	if err != nil {
+		return nil, fmt.Errorf("query rows: %w", err)
+	}
+
+	return func(yield func(PT, error) bool) {
+		defer rows.Close()
+		for rows.Next() {
+			var record T
+			if err := t.Query.Scan.ScanOne(&record, rows); err != nil {
+				if !errors.Is(err, pgx.ErrNoRows) {
+					yield(nil, err)
+				}
+				return
+			}
+			if !yield(&record, nil) {
+				return
+			}
+		}
+	}, nil
+}
+
 // GetByID returns a record by its ID.
 func (t *Table[T, PT, IDT]) GetByID(ctx context.Context, id IDT) (PT, error) {
 	return t.Get(ctx, sq.Eq{t.IDColumn: id}, []string{t.IDColumn})

tests/table_test.go

@@ -54,6 +54,15 @@ func TestTable(t *testing.T) {
 		count, err = db.Accounts.Count(ctx, nil)
 		require.NoError(t, err, "FindAll failed")
 		require.Equal(t, uint64(1), count, "Expected 1 account")
+
+		// Iterate all accounts.
+		iter, err := db.Accounts.Iter(ctx, nil, nil)
+		require.NoError(t, err, "Iter failed")
+		var accounts []Account
+		for account, err := range iter {
+			require.NoError(t, err, "Iter error")
+			accounts = append(accounts, *account)
+		}
 	})
 
 	t.Run("Save multiple", func(t *testing.T) {