Implement Batch

batch.go

@@ -0,0 +1,65 @@
+package sequel
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"slices"
+	"strings"
+)
+
+// Executor is the interface for types that can execute queries. Both [DB] and
+// [Tx] satisfy this interface.
+type Executor interface {
+	ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
+}
+
+// BatchSize is the maximum number of records inserted per statement.
+const BatchSize = 100
+
+// Batch inserts a slice of items into the given table using multi-row INSERT
+// statements. Items are inserted in chunks of [BatchSize]. The columns
+// parameter specifies the column names, and the values function maps each item to
+// its column values. The length of the slice returned by extractValues must match the
+// length of columns. Batch does nothing if items is empty. Table, columns and onConfict
+// are not sanitized; they must come from a trusted source. The extractValues function will
+// never be called concurrently.
+func Batch[T any](ctx context.Context, exec Executor, table string, columns []string, onConflict string, items []T, extractValues func(T) []any) error {
+	batch := 0
+	for chunk := range slices.Chunk(items, BatchSize) {
+		query, args := batchQuery(table, columns, onConflict, chunk, extractValues)
+		if _, err := exec.ExecContext(ctx, query, args...); err != nil {
+			return fmt.Errorf("batch %d (%d items) failed: %w", batch, len(chunk), err)
+		}
+		batch++
+	}
+	return nil
+}
+
+func batchQuery[T any](table string, columns []string, onConflict string, items []T, extractValues func(T) []any) (string, []any) {
+	ncols := len(columns)
+	args := make([]any, 0, len(items)*ncols)
+
+	var b strings.Builder
+	fmt.Fprintf(&b, "INSERT INTO %s (%s) VALUES ", table, strings.Join(columns, ", "))
+
+	for i, item := range items {
+		if i > 0 {
+			b.WriteString(", ")
+		}
+		b.WriteByte('(')
+		vals := extractValues(item)
+		for j, v := range vals {
+			if j > 0 {
+				b.WriteString(", ")
+			}
+			fmt.Fprintf(&b, "$%d", i*ncols+j+1)
+			args = append(args, v)
+		}
+		b.WriteByte(')')
+	}
+
+	fmt.Fprintf(&b, " %s", onConflict)
+
+	return b.String(), args
+}

batch_test.go

@@ -0,0 +1,136 @@
+package sequel
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBatch(t *testing.T) {
+	db, err := New(postgresDataSource)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		assert.NoError(t, db.Close())
+	})
+
+	ctx := context.Background()
+	t.Cleanup(func() {
+		_, _ = db.Exec(ctx, "DELETE FROM person_test WHERE name LIKE 'batch-%'")
+	})
+
+	type person struct {
+		Name  string
+		Email string
+	}
+
+	columns := []string{"name", "email"}
+	fn := func(p person) []any {
+		return []any{p.Name, p.Email}
+	}
+
+	t.Run("empty", func(t *testing.T) {
+		err := Batch(ctx, nil, "person_test", columns, "", []person{}, fn)
+		assert.NoError(t, err)
+	})
+
+	t.Run("single", func(t *testing.T) {
+		items := []person{
+			{Name: "batch-single", Email: "batch-single@example.com"},
+		}
+		require.NoError(t, Batch(ctx, db, "person_test", columns, "", items, fn))
+
+		var name string
+		err := db.QueryRow(ctx, "SELECT name FROM person_test WHERE email = $1", "batch-single@example.com").Scan(&name)
+		require.NoError(t, err)
+		assert.Equal(t, "batch-single", name)
+	})
+
+	t.Run("multiple", func(t *testing.T) {
+		items := []person{
+			{Name: "batch-multi-1", Email: "batch-multi-1@example.com"},
+			{Name: "batch-multi-2", Email: "batch-multi-2@example.com"},
+			{Name: "batch-multi-3", Email: "batch-multi-3@example.com"},
+		}
+		require.NoError(t, Batch(ctx, db, "person_test", columns, "", items, fn))
+
+		for _, p := range items {
+			var name string
+			err := db.QueryRow(ctx, "SELECT name FROM person_test WHERE email = $1", p.Email).Scan(&name)
+			require.NoError(t, err)
+			assert.Equal(t, p.Name, name)
+		}
+	})
+
+	t.Run("chunked", func(t *testing.T) {
+		// Insert more than BatchSize to verify chunking.
+		items := make([]person, BatchSize+5)
+		for i := range items {
+			items[i] = person{
+				Name:  fmt.Sprintf("batch-chunk-%d", i),
+				Email: fmt.Sprintf("batch-chunk-%d@example.com", i),
+			}
+		}
+		require.NoError(t, Batch(ctx, db, "person_test", columns, "", items, fn))
+
+		var count int
+		err := db.QueryRow(ctx, "SELECT COUNT(*) FROM person_test WHERE name LIKE 'batch-chunk-%'").Scan(&count)
+		require.NoError(t, err)
+		assert.Equal(t, len(items), count)
+	})
+
+	t.Run("tx", func(t *testing.T) {
+		tx, err := db.Begin(ctx)
+		require.NoError(t, err)
+
+		items := []person{
+			{Name: "batch-tx-1", Email: "batch-tx-1@example.com"},
+			{Name: "batch-tx-2", Email: "batch-tx-2@example.com"},
+		}
+		require.NoError(t, Batch(ctx, tx, "person_test", columns, "", items, fn))
+		require.NoError(t, tx.Commit())
+
+		var count int
+		err = db.QueryRow(ctx, "SELECT COUNT(*) FROM person_test WHERE name LIKE 'batch-tx-%'").Scan(&count)
+		require.NoError(t, err)
+		assert.Equal(t, 2, count)
+	})
+
+	t.Run("onConflictDoNothing", func(t *testing.T) {
+		items := []person{
+			{Name: "batch-conflict-1", Email: "batch-conflict-1@example.com"},
+		}
+		require.NoError(t, Batch(ctx, db, "person_test", columns, "", items, fn))
+
+		// Insert again with a different name but same email; conflict should be ignored.
+		dupes := []person{
+			{Name: "batch-conflict-1-updated", Email: "batch-conflict-1@example.com"},
+		}
+		require.NoError(t, Batch(ctx, db, "person_test", columns, "ON CONFLICT (email) DO NOTHING", dupes, fn))
+
+		var name string
+		err := db.QueryRow(ctx, "SELECT name FROM person_test WHERE email = $1", "batch-conflict-1@example.com").Scan(&name)
+		require.NoError(t, err)
+		assert.Equal(t, "batch-conflict-1", name)
+	})
+
+	t.Run("onConflictDoUpdate", func(t *testing.T) {
+		items := []person{
+			{Name: "batch-upsert-1", Email: "batch-upsert-1@example.com"},
+		}
+		require.NoError(t, Batch(ctx, db, "person_test", columns, "", items, fn))
+
+		// Insert again with a different name but same email; name should be updated.
+		upserts := []person{
+			{Name: "batch-upsert-1-updated", Email: "batch-upsert-1@example.com"},
+		}
+		require.NoError(t, Batch(ctx, db, "person_test", columns, "ON CONFLICT (email) DO UPDATE SET name = EXCLUDED.name", upserts, fn))
+
+		var name string
+		err := db.QueryRow(ctx, "SELECT name FROM person_test WHERE email = $1", "batch-upsert-1@example.com").Scan(&name)
+		require.NoError(t, err)
+		assert.Equal(t, "batch-upsert-1-updated", name)
+	})
+}

sequel.go

@@ -290,6 +290,12 @@ func (d *DB) Exec(ctx context.Context, query string, args ...any) (sql.Result, e
 	return d.db.ExecContext(ctx, query, args...)
 }
 
+// ExecContext is an alias for [DB.Exec] that satisfies the [Executor]
+// interface.
+func (d *DB) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
+	return d.db.ExecContext(ctx, query, args...)
+}
+
 // RebindQuery executes a query that returns rows, typically a SELECT. The query is
 // rebound from `?` to the DB driver's bind type. The args are for any
 // placeholder parameters in the query.