Add CLI statement-timeout flag

River CLI Postgres commands had fixed session defaults and no explicit
flag for statement timeout. For long-running migrations, callers either
accepted the default or had to encode `statement_timeout` into the
connection string. This could be difficult in deployed environments
where i.e. a database connection string was shared with an application
and the user did not want to accidentally affect timeouts for other
processes using a shared connection string.

This commit adds a root `--statement-timeout` flag and thread it through
command bootstrapping into Postgres pool configuration. Statement
timeout now uses explicit precedence: flag value first, then
connection-string query parameter, then built-in default.
`idle_in_transaction_session_timeout` keeps its existing behavior and
remains URL/default driven.

Author: bgentry

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Added root River CLI flag `--statement-timeout` so Postgres session statement timeout can be set explicitly for commands like migrations. Explicit flag values take priority over database URL query params, and query params still take priority over built-in defaults. [PR #XXXX](https://github.com/riverqueue/river/pull/XXXX).
+
 ### Fixed
 
 - `JobCountByQueueAndState` now returns consistent results across drivers, including requested queues with zero jobs, and deduplicates repeated queue names in input. This resolves an issue with the sqlite driver in River UI reported in [riverqueue/riverui#496](https://github.com/riverqueue/riverui#496). [PR #1140](https://github.com/riverqueue/river/pull/1140).

cmd/river/rivercli/command.go

@@ -1,6 +1,7 @@
 package rivercli
 
 import (
+	"cmp"
 	"context"
 	"database/sql"
 	"fmt"
@@ -46,11 +47,12 @@ type CommandOpts interface {
 
 // RunCommandBundle is a bundle of utilities for RunCommand.
 type RunCommandBundle struct {
-	DatabaseURL    *string
-	DriverProcurer DriverProcurer
-	Logger         *slog.Logger
-	OutStd         io.Writer
-	Schema         string
+	DatabaseURL      *string
+	DriverProcurer   DriverProcurer
+	Logger           *slog.Logger
+	OutStd           io.Writer
+	Schema           string
+	StatementTimeout *time.Duration
 }
 
 // RunCommand bootstraps and runs a River CLI subcommand.
@@ -81,7 +83,7 @@ func RunCommand[TOpts CommandOpts](ctx context.Context, bundle *RunCommandBundle
 		if databaseURL != nil {
 			switch protocol {
 			case "postgres", "postgresql":
-				dbPool, err := openPgxV5DBPool(ctx, *databaseURL)
+				dbPool, err := openPgxV5DBPool(ctx, *databaseURL, bundle.StatementTimeout)
 				if err != nil {
 					return false, err
 				}
@@ -128,7 +130,7 @@ func RunCommand[TOpts CommandOpts](ctx context.Context, bundle *RunCommandBundle
 	return nil
 }
 
-func openPgxV5DBPool(ctx context.Context, databaseURL string) (*pgxpool.Pool, error) {
+func openPgxV5DBPool(ctx context.Context, databaseURL string, statementTimeout *time.Duration) (*pgxpool.Pool, error) {
 	const (
 		defaultIdleInTransactionSessionTimeout = 11 * time.Second // should be greater than statement timeout because statements count towards idle-in-transaction
 		defaultStatementTimeout                = 10 * time.Second
@@ -149,9 +151,27 @@ func openPgxV5DBPool(ctx context.Context, databaseURL string) (*pgxpool.Pool, er
 		runtimeParams[name] = val
 	}
 
-	setParamIfUnset(pgxConfig.ConnConfig.RuntimeParams, "application_name", "river CLI")
-	setParamIfUnset(pgxConfig.ConnConfig.RuntimeParams, "idle_in_transaction_session_timeout", strconv.Itoa(int(defaultIdleInTransactionSessionTimeout.Milliseconds())))
-	setParamIfUnset(pgxConfig.ConnConfig.RuntimeParams, "statement_timeout", strconv.Itoa(int(defaultStatementTimeout.Milliseconds())))
+	millisecondsStringFromDurationPtr := func(duration *time.Duration) string {
+		if duration == nil {
+			return ""
+		}
+
+		return strconv.Itoa(int(duration.Milliseconds()))
+	}
+
+	runtimeParams := pgxConfig.ConnConfig.RuntimeParams
+	if runtimeParams == nil {
+		runtimeParams = make(map[string]string)
+		pgxConfig.ConnConfig.RuntimeParams = runtimeParams
+	}
+
+	setParamIfUnset(runtimeParams, "application_name", "river CLI")
+	setParamIfUnset(runtimeParams, "idle_in_transaction_session_timeout", strconv.Itoa(int(defaultIdleInTransactionSessionTimeout.Milliseconds())))
+	runtimeParams["statement_timeout"] = cmp.Or(
+		millisecondsStringFromDurationPtr(statementTimeout),
+		runtimeParams["statement_timeout"],
+		strconv.Itoa(int(defaultStatementTimeout.Milliseconds())),
+	)
 
 	dbPool, err := pgxpool.NewWithConfig(ctx, pgxConfig)
 	if err != nil {

cmd/river/rivercli/river_cli.go

@@ -60,9 +60,11 @@ func (c *CLI) BaseCommandSet() *cobra.Command {
 	ctx := context.Background()
 
 	var globalOpts struct {
-		Debug   bool
-		Verbose bool
+		Debug            bool
+		StatementTimeout time.Duration
+		Verbose          bool
 	}
+	var rootCmd *cobra.Command
 
 	makeLogger := func() *slog.Logger {
 		switch {
@@ -77,16 +79,21 @@ func (c *CLI) BaseCommandSet() *cobra.Command {
 
 	// Make a bundle for RunCommand. Takes a database URL pointer because not every command is required to take a database URL.
 	makeCommandBundle := func(databaseURL *string, schema string) *RunCommandBundle {
+		var statementTimeout *time.Duration
+		if rootCmd.PersistentFlags().Changed("statement-timeout") {
+			statementTimeout = &globalOpts.StatementTimeout
+		}
+
 		return &RunCommandBundle{
-			DatabaseURL:    databaseURL,
-			DriverProcurer: c.driverProcurer,
-			Logger:         makeLogger(),
-			OutStd:         c.out,
-			Schema:         schema,
+			DatabaseURL:      databaseURL,
+			DriverProcurer:   c.driverProcurer,
+			Logger:           makeLogger(),
+			OutStd:           c.out,
+			Schema:           schema,
+			StatementTimeout: statementTimeout,
 		}
 	}
 
-	var rootCmd *cobra.Command
 	{
 		var rootOpts struct {
 			Version bool
@@ -116,6 +123,7 @@ PG* vars if it's been specified.
 		rootCmd.SetOut(c.out)
 
 		rootCmd.PersistentFlags().BoolVar(&globalOpts.Debug, "debug", false, "output maximum logging verbosity (debug level)")
+		rootCmd.PersistentFlags().DurationVar(&globalOpts.StatementTimeout, "statement-timeout", 0, "override Postgres statement_timeout (Go duration, e.g. 10s, 1m)")
 		rootCmd.PersistentFlags().BoolVarP(&globalOpts.Verbose, "verbose", "v", false, "output additional logging verbosity (info level)")
 		rootCmd.MarkFlagsMutuallyExclusive("debug", "verbose")
 
@@ -151,7 +159,7 @@ before starting the client, and works until all jobs are finished.
 
 The database in --database-url will have its jobs table truncated, so make sure
 to use a development database only.
-	`),
+		`),
 			RunE: func(cmd *cobra.Command, args []string) error {
 				return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL, opts.Schema), &bench{}, &opts)
 			},
@@ -192,7 +200,7 @@ Defaults to running a single down migration. This behavior can be changed with
 SQL being run can be output using --show-sql, and executing real database
 operations can be prevented with --dry-run. Combine --show-sql and --dry-run to
 dump prospective migrations that would be applied to stdout.
-	`),
+		`),
 			RunE: func(cmd *cobra.Command, args []string) error {
 				return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL, opts.Schema), &migrateDown{}, &opts)
 			},
@@ -235,7 +243,7 @@ Returns migrations for Postgres by default. Use --database-url to hint a
 non-Postgres alternative:
 
     river migrate-get --version 6 --up --database-url sqlite:// > river_6.up.sql
-	`),
+		`),
 			RunE: func(cmd *cobra.Command, args []string) error {
 				return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL, ""), &migrateGet{}, &opts)
 			},
@@ -265,7 +273,7 @@ non-Postgres alternative:
 			Long: strings.TrimSpace(`
 List available migrations for the given line, showing the currently applied
 migration as determined by the database in --database-url.
-	`),
+		`),
 			RunE: func(cmd *cobra.Command, args []string) error {
 				return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL, opts.Schema), &migrateList{}, &opts)
 			},
@@ -292,7 +300,7 @@ restricted with --max-steps or --target-version.
 SQL being run can be output using --show-sql, and executing real database
 operations can be prevented with --dry-run. Combine --show-sql and --dry-run to
 dump prospective migrations that would be applied to stdout.
-	`),
+		`),
 			RunE: func(cmd *cobra.Command, args []string) error {
 				return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL, opts.Schema), &migrateUp{}, &opts)
 			},
@@ -314,7 +322,7 @@ are outstanding migrations that still need to be run.
 
 Can be paired with river migrate-up --dry-run --show-sql to dump information on
 migrations that need to be run, but without running them.
-	`),
+		`),
 			RunE: func(cmd *cobra.Command, args []string) error {
 				return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL, opts.Schema), &validate{}, &opts)
 			},
@@ -332,7 +340,7 @@ migrations that need to be run, but without running them.
 			Short: "Print version information",
 			Long: strings.TrimSpace(`
 Print River and Go version information.
-	`),
+		`),
 			RunE: func(cmd *cobra.Command, args []string) error {
 				return RunCommand(ctx, makeCommandBundle(nil, ""), &version{}, &versionOpts{Name: c.name})
 			},

cmd/river/rivercli/river_cli_test.go

@@ -5,6 +5,7 @@ import (
 	"cmp"
 	"context"
 	"fmt"
+	"maps"
 	"net/url"
 	"runtime/debug"
 	"strings"
@@ -261,6 +262,109 @@ func TestBaseCommandSetNonParallel(t *testing.T) {
 		cmd.SetArgs([]string{"migrate-up", "--schema", schema})
 		require.NoError(t, cmd.Execute())
 	})
+
+	t.Run("PostgresTimeoutPrecedence", func(t *testing.T) {
+		type testCase struct {
+			args                       []string
+			expectedStatementTimeoutMS string
+			name                       string
+		}
+
+		makeCommandAndParams := func(t *testing.T) (*cobra.Command, func() map[string]string) {
+			t.Helper()
+
+			var capturedRuntimeParams map[string]string
+
+			migratorStub := &MigratorStub{}
+			migratorStub.allVersionsStub = func() []rivermigrate.Migration { return []rivermigrate.Migration{testMigration01} }
+			migratorStub.getVersionStub = func(version int) (rivermigrate.Migration, error) {
+				if version == 1 {
+					return testMigration01, nil
+				}
+
+				return rivermigrate.Migration{}, fmt.Errorf("unknown version: %d", version)
+			}
+			migratorStub.existingVersionsStub = func(ctx context.Context) ([]rivermigrate.Migration, error) { return nil, nil }
+
+			cli := NewCLI(&Config{
+				DriverProcurer: &DriverProcurerStub{
+					getMigratorStub: func(config *rivermigrate.Config) (MigratorInterface, error) {
+						return migratorStub, nil
+					},
+					initPgxV5Stub: func(pool *pgxpool.Pool) {
+						capturedRuntimeParams = maps.Clone(pool.Config().ConnConfig.RuntimeParams)
+					},
+				},
+				Name: "River",
+			})
+
+			var out bytes.Buffer
+			cli.SetOut(&out)
+
+			return cli.BaseCommandSet(), func() map[string]string {
+				return capturedRuntimeParams
+			}
+		}
+
+		makeBaseDatabaseURL := func(t *testing.T) *url.URL {
+			t.Helper()
+
+			testDatabaseURL := riversharedtest.TestDatabaseURL()
+			parsedDatabaseURL, err := url.Parse(testDatabaseURL)
+			require.NoError(t, err)
+
+			return parsedDatabaseURL
+		}
+
+		testCases := []testCase{
+			{
+				name:                       "DefaultsAppliedWhenNothingSpecified",
+				args:                       []string{},
+				expectedStatementTimeoutMS: "10000",
+			},
+			{
+				name:                       "DatabaseURLQueryParamsOverrideDefaults",
+				args:                       []string{},
+				expectedStatementTimeoutMS: "11234",
+			},
+			{
+				name: "ExplicitFlagsOverrideDatabaseURLQueryParams",
+				args: []string{
+					"--statement-timeout", "1m3.123s",
+				},
+				expectedStatementTimeoutMS: "63123",
+			},
+		}
+
+		for _, tc := range testCases {
+			t.Run(tc.name, func(t *testing.T) {
+				cmd, getRuntimeParams := makeCommandAndParams(t)
+
+				databaseURL := makeBaseDatabaseURL(t)
+				switch tc.name {
+				case "DatabaseURLQueryParamsOverrideDefaults":
+					queryValues := databaseURL.Query()
+					queryValues.Set("statement_timeout", "11234")
+					databaseURL.RawQuery = queryValues.Encode()
+				case "ExplicitFlagsOverrideDatabaseURLQueryParams":
+					queryValues := databaseURL.Query()
+					queryValues.Set("statement_timeout", "12345")
+					databaseURL.RawQuery = queryValues.Encode()
+				}
+
+				args := append([]string{
+					"migrate-get", "--up", "--version", "1", "--database-url", databaseURL.String(),
+				}, tc.args...)
+				cmd.SetArgs(args)
+				require.NoError(t, cmd.Execute())
+
+				runtimeParams := getRuntimeParams()
+				require.NotNil(t, runtimeParams)
+
+				require.Equal(t, tc.expectedStatementTimeoutMS, runtimeParams["statement_timeout"])
+			})
+		}
+	})
 }
 
 func TestBaseCommandSetDriverProcurerPgxV5(t *testing.T) {

docs/development.md

@@ -30,6 +30,11 @@ To run programs locally outside of tests, create and raise a development databas
     createdb river_dev
     go run ./cmd/river migrate-up --database-url postgres:///river_dev --line main
 
+If needed, override Postgres timeouts for long-running migrations with root CLI
+flags:
+
+    go run ./cmd/river --statement-timeout 2m migrate-up --database-url postgres:///river_dev --line main
+
 ## Releasing a new version
 
 1. Fetch changes to the repo and any new tags. Export `VERSION` by incrementing the last tag. Execute `update-mod-version` to add it the project's `go.mod` files:
@@ -71,4 +76,4 @@ Modify `go.work` so it contains the new desired version in `go` and/or `toolchai
 
 ```shell
 make update-mod-go
-```
+```