diff --git a/CHANGELOG.md b/CHANGELOG.md index 62af72d0..8f3a8bef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Added hook for River's CLI command framework that can be used to inject custom SQLite initialization. Used by River Pro. [PR #1369](https://github.com/riverqueue/river/pull/1369). + ## [0.46.0] - 2026-08-29 ### Added diff --git a/cmd/river/rivercli/command.go b/cmd/river/rivercli/command.go index 9bdce036..317e9607 100644 --- a/cmd/river/rivercli/command.go +++ b/cmd/river/rivercli/command.go @@ -101,7 +101,12 @@ func RunCommand[TOpts CommandOpts](ctx context.Context, bundle *RunCommandBundle } defer dbPool.Close() - driverProcurer = &sqliteDriverProcurer{dbPool: dbPool} + driverProcurerSQLite, isSQLiteProcurer := driverProcurer.(DriverProcurerSQLite) + if driverProcurer != nil && isSQLiteProcurer { + driverProcurerSQLite.InitSQLite(dbPool) + } else { + driverProcurer = &sqliteDriverProcurer{dbPool: dbPool} + } default: return false, fmt.Errorf("unsupported database URL (`%s`); try one with a `postgres://`, `postgresql://`, or `sqlite://` scheme/prefix", *bundle.DatabaseURL) diff --git a/cmd/river/rivercli/driver_procurer.go b/cmd/river/rivercli/driver_procurer.go index 29e4edce..181b4317 100644 --- a/cmd/river/rivercli/driver_procurer.go +++ b/cmd/river/rivercli/driver_procurer.go @@ -23,10 +23,16 @@ type DriverProcurer interface { QueryRow(ctx context.Context, sql string, args ...any) riverdriver.Row } +// DriverProcurerPgxV5 optionally initializes a custom procurer with a pgx/v5 pool. type DriverProcurerPgxV5 interface { InitPgxV5(pool *pgxpool.Pool) } +// DriverProcurerSQLite optionally initializes a custom procurer with a SQLite pool. +type DriverProcurerSQLite interface { + InitSQLite(pool *sql.DB) +} + // BenchmarkerInterface is an interface to a Benchmarker. Its reason for // existence is to wrap a benchmarker to strip it of its generic parameter, // letting us pass it around without having to know the transaction type. diff --git a/cmd/river/rivercli/river_cli_test.go b/cmd/river/rivercli/river_cli_test.go index 95ee4fe5..debeef9b 100644 --- a/cmd/river/rivercli/river_cli_test.go +++ b/cmd/river/rivercli/river_cli_test.go @@ -4,6 +4,7 @@ import ( "bytes" "cmp" "context" + "database/sql" "fmt" "maps" "net/url" @@ -28,6 +29,7 @@ type DriverProcurerStub struct { getBenchmarkerStub func(config *riverbench.Config) BenchmarkerInterface getMigratorStub func(config *rivermigrate.Config) (MigratorInterface, error) initPgxV5Stub func(pool *pgxpool.Pool) + initSQLiteStub func(pool *sql.DB) queryRowStub func(ctx context.Context, sql string, args ...any) riverdriver.Row } @@ -55,6 +57,14 @@ func (p *DriverProcurerStub) InitPgxV5(pool *pgxpool.Pool) { p.initPgxV5Stub(pool) } +func (p *DriverProcurerStub) InitSQLite(pool *sql.DB) { + if p.initSQLiteStub == nil { + panic("InitSQLite is not stubbed") + } + + p.initSQLiteStub(pool) +} + func (p *DriverProcurerStub) QueryRow(ctx context.Context, sql string, args ...any) riverdriver.Row { if p.queryRowStub == nil { panic("QueryRow is not stubbed") @@ -447,6 +457,52 @@ SELECT 'up 1' FROM river_table `), strings.TrimSpace(out.String())) } +func TestBaseCommandSetDriverProcurerSQLite(t *testing.T) { + t.Parallel() + + getMigratorCalled := false + initSQLiteCalled := false + + 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) { + getMigratorCalled = true + return migratorStub, nil + }, + initSQLiteStub: func(pool *sql.DB) { + initSQLiteCalled = true + }, + }, + Name: "River", + }) + + var out bytes.Buffer + cli.SetOut(&out) + + cmd := cli.BaseCommandSet() + cmd.SetArgs([]string{"migrate-get", "--up", "--version", "1", "--database-url", "sqlite://"}) + require.NoError(t, cmd.Execute()) + + require.True(t, getMigratorCalled) + require.True(t, initSQLiteCalled) + + require.Equal(t, strings.TrimSpace(` +-- River main migration 001 [up] +SELECT 'up 1' FROM river_table + `), strings.TrimSpace(out.String())) +} + func TestMigrateGet(t *testing.T) { t.Parallel()