Browse Source

Add Each and EachPrepared methods to Tx

Volodymyr Tkach 2 years ago
parent
commit
892d72f01c
1 changed files with 30 additions and 0 deletions
  1. 30 0
      gosql/common/tx.go

+ 30 - 0
gosql/common/tx.go

@@ -3,6 +3,7 @@ package common
 import (
 	"context"
 	"database/sql"
+	"fmt"
 	"os"
 
 	"time"
@@ -45,6 +46,35 @@ func (t *Tx) DeleteRowByID(ctx context.Context, id int64, row any) error {
 	return err
 }
 
+func (t *Tx) Each(ctx context.Context, query string, callback func(ctx context.Context, rows *Rows) error, args ...any) error {
+	if callback == nil {
+		return fmt.Errorf("callback is not set")
+	}
+	rows, err := t.Query(ctx, query, args...)
+	if err != nil {
+		return err
+	}
+	defer rows.Close()
+	for rows.Next() {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			if err := callback(ctx, rows); err != nil {
+				return err
+			}
+		}
+	}
+	if err := rows.Err(); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (t *Tx) EachPrepared(ctx context.Context, prep *Prepared, callback func(ctx context.Context, rows *Rows) error) error {
+	return t.Each(ctx, prep.Query, callback, prep.Args...)
+}
+
 func (t *Tx) Exec(ctx context.Context, query string, args ...any) (sql.Result, error) {
 	start := time.Now()
 	res, err := t.tx.ExecContext(ctx, t.fixQuery(query), args...)