crdb: Add ExecuteCtx for configurable retry limits (#180)

* crdb: Add ExecuteCtx for configurable retry limits

Adds ExecuteCtx which allows configuring maximum retries through context,
alongside the existing Execute function. This enables better control over
retry behavior while maintaining compatibility with code using the original
Execute.

The new function helps prevent infinite retry loops in environments where
bounded retry counts are required, while preserving the existing error
handling and retry semantics.

Author: sean-

.github/workflows/ci.yaml

@@ -27,10 +27,7 @@ jobs:
     strategy:
       matrix:
         go:
-          - "1.19"
-          - "1.20"
-          - "1.21"
-          - "1.22"
+          - "1.23"
 
     steps:
     - uses: actions/checkout@v2
@@ -42,11 +39,9 @@ jobs:
 
     - name: Install dependencies
       run: |
-        # The latest crlfmt requires go1.19.
-        go install github.com/cockroachdb/[email protected]
-        go install github.com/kisielk/[email protected]
+        go install github.com/kisielk/errcheck@latest
         go install github.com/mdempsky/unconvert@latest
-        go install honnef.co/go/tools/cmd/staticcheck@2023.1.7
+        go install honnef.co/go/tools/cmd/staticcheck@latest
         go install golang.org/x/tools/go/analysis/passes/shadow/cmd/shadow@latest
 
     - name: Build

crdb/tx.go

@@ -100,6 +100,64 @@ func Execute(fn func() error) (err error) {
 	}
 }
 
+// ExecuteCtxFunc represents a function that takes a context and variadic
+// arguments and returns an error. It's used with ExecuteCtx to enable retryable
+// operations with configurable parameters.
+type ExecuteCtxFunc func(context.Context, ...interface{}) error
+
+// ExecuteCtx runs fn and retries it as needed, respecting a maximum retry count
+// obtained from the context. It is used to add configurable retry handling to
+// the execution of a single statement. If a multi-statement transaction is
+// being run, use ExecuteTx instead.
+//
+// The maximum number of retries can be configured using WithMaxRetries(ctx, n).
+// Setting n=0 allows one attempt with no retries. If the number of retries is
+// exhausted and the last attempt resulted in a retryable error, ExecuteCtx
+// returns a max retries exceeded error wrapping the last retryable error
+// encountered.
+//
+// The fn parameter accepts variadic arguments which are passed through on each
+// retry attempt, allowing for flexible parameterization of the retried operation.
+//
+// As with Execute, retry handling for individual statements (implicit transactions)
+// is usually performed automatically on the CockroachDB SQL gateway, making use
+// of this function generally unnecessary. However, automatic retries are disabled
+// once result streaming begins (typically when results exceed 16KiB).
+//
+// NOTE: the supplied fn closure should not have external side effects beyond
+// changes to the database.
+//
+// fn must take care when wrapping errors returned from the database driver with
+// additional context. To preserve retry behavior, errors should implement either
+// `Cause() error` (github.com/pkg/errors) or `Unwrap() error` (Go 1.13+).
+// For example:
+//
+//	crdb.ExecuteCtx(ctx, func(ctx context.Context, args ...interface{}) error {
+//	    id := args[0].(int)
+//	    rows, err := db.QueryContext(ctx, "SELECT * FROM users WHERE id = $1", id)
+//	    if err != nil {
+//	        return fmt.Errorf("scanning row: %w", err) // uses %w for proper error wrapping
+//	    }
+//	    defer rows.Close()
+//	    // ...
+//	    return nil
+//	}, userID)
+func ExecuteCtx(ctx context.Context, fn ExecuteCtxFunc, args ...interface{}) (err error) {
+	maxRetries := numRetriesFromContext(ctx)
+	for n := 0; n <= maxRetries; n++ {
+		if err = ctx.Err(); err != nil {
+			return err
+		}
+
+		err = fn(ctx, args...)
+		if err == nil || !errIsRetryable(err) {
+			return err
+		}
+	}
+
+	return newMaxRetriesExceededError(err, maxRetries)
+}
+
 type txConfigKey struct{}
 
 // WithMaxRetries configures context so that ExecuteTx retries tx specified

crdb/tx_test.go

@@ -17,12 +17,75 @@ package crdb
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"fmt"
 	"testing"
 
 	"github.com/cockroachdb/cockroach-go/v2/testserver"
 )
 
+// TestExecuteCtx verifies that ExecuteCtx correctly handles different retry limits
+// and context cancellation when executing database operations.
+//
+// TODO(seanc@): Add test cases that force retryable errors by simulating
+// transaction conflicts or network failures. Consider using the same write skew
+// pattern from TestExecuteTx.
+func TestExecuteCtx(t *testing.T) {
+	db, stop := testserver.NewDBForTest(t)
+	defer stop()
+	ctx := context.Background()
+
+	// Setup test table
+	if _, err := db.ExecContext(ctx, `CREATE TABLE test_retry (id INT PRIMARY KEY)`); err != nil {
+		t.Fatal(err)
+	}
+
+	testCases := []struct {
+		name       string
+		maxRetries int
+		id         int
+		withCancel bool
+		wantErr    error
+	}{
+		{"no retries", 0, 0, false, nil},
+		{"single retry", 1, 1, false, nil},
+		{"cancelled context", 1, 2, true, context.Canceled},
+		{"no args", 1, 3, false, nil},
+	}
+
+	fn := func(ctx context.Context, args ...interface{}) error {
+		if len(args) == 0 {
+			_, err := db.ExecContext(ctx, `INSERT INTO test_retry VALUES (3)`)
+			return err
+		}
+		id := args[0].(int)
+		_, err := db.ExecContext(ctx, `INSERT INTO test_retry VALUES ($1)`, id)
+		return err
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			limitedCtx := WithMaxRetries(ctx, tc.maxRetries)
+			if tc.withCancel {
+				var cancel context.CancelFunc
+				limitedCtx, cancel = context.WithCancel(limitedCtx)
+				cancel()
+			}
+
+			var err error
+			if tc.name == "no args" {
+				err = ExecuteCtx(limitedCtx, fn)
+			} else {
+				err = ExecuteCtx(limitedCtx, fn, tc.id)
+			}
+
+			if !errors.Is(err, tc.wantErr) {
+				t.Errorf("got error %v, want %v", err, tc.wantErr)
+			}
+		})
+	}
+}
+
 // TestExecuteTx verifies transaction retry using the classic
 // example of write skew in bank account balance transfers.
 func TestExecuteTx(t *testing.T) {

go.mod

@@ -1,17 +1,17 @@
 module github.com/cockroachdb/cockroach-go/v2
 
-go 1.19
+go 1.23
 
 require (
-	github.com/gofrs/flock v0.8.1
+	github.com/gofrs/flock v0.12.1
 	github.com/jackc/pgx/v4 v4.18.3
-	github.com/jackc/pgx/v5 v5.5.2
-	github.com/jmoiron/sqlx v1.3.5
-	github.com/lib/pq v1.10.6
-	github.com/stretchr/testify v1.8.1
+	github.com/jackc/pgx/v5 v5.7.2
+	github.com/jmoiron/sqlx v1.4.0
+	github.com/lib/pq v1.10.9
+	github.com/stretchr/testify v1.10.0
 	gopkg.in/yaml.v3 v3.0.1
-	gorm.io/driver/postgres v1.3.5
-	gorm.io/gorm v1.23.5
+	gorm.io/driver/postgres v1.5.11
+	gorm.io/gorm v1.25.12
 )
 
 require (
@@ -21,16 +21,16 @@ require (
 	github.com/jackc/pgio v1.0.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgproto3/v2 v2.3.3 // indirect
-	github.com/jackc/pgservicefile v0.0.0-20231201235250-de7065d80cb9 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
 	github.com/jackc/pgtype v1.14.3 // indirect
 	github.com/jackc/puddle v1.3.0 // indirect
-	github.com/jackc/puddle/v2 v2.2.1 // indirect
+	github.com/jackc/puddle/v2 v2.2.2 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
-	github.com/jinzhu/now v1.1.4 // indirect
+	github.com/jinzhu/now v1.1.5 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rogpeppe/go-internal v1.9.0 // indirect
-	golang.org/x/crypto v0.22.0 // indirect
-	golang.org/x/sync v0.7.0 // indirect
-	golang.org/x/sys v0.19.0 // indirect
-	golang.org/x/text v0.14.0 // indirect
+	golang.org/x/crypto v0.31.0 // indirect
+	golang.org/x/sync v0.10.0 // indirect
+	golang.org/x/sys v0.28.0 // indirect
+	golang.org/x/text v0.21.0 // indirect
 )