Skip to content

Commit caffc57

Browse files
committed
Add hook to support Pro SQLite driver
Here, add a custom hook to SQLite driver procurement that can be used by River Pro to inject some needed customization. A similar hook already exists for Pgx V5. This was actually needed in v0.30.0 to work properly, so I'll have to put in a patch release.
1 parent 195adeb commit caffc57

4 files changed

Lines changed: 72 additions & 1 deletion

File tree

‎CHANGELOG.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
77

88
## [Unreleased]
99

10+
### Added
11+
12+
- 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).
13+
1014
## [0.46.0] - 2026-08-29
1115

1216
### Added

‎cmd/river/rivercli/command.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,12 @@ func RunCommand[TOpts CommandOpts](ctx context.Context, bundle *RunCommandBundle
101101
}
102102
defer dbPool.Close()
103103

104-
driverProcurer = &sqliteDriverProcurer{dbPool: dbPool}
104+
driverProcurerSQLite, isSQLiteProcurer := driverProcurer.(DriverProcurerSQLite)
105+
if driverProcurer != nil && isSQLiteProcurer {
106+
driverProcurerSQLite.InitSQLite(dbPool)
107+
} else {
108+
driverProcurer = &sqliteDriverProcurer{dbPool: dbPool}
109+
}
105110

106111
default:
107112
return false, fmt.Errorf("unsupported database URL (`%s`); try one with a `postgres://`, `postgresql://`, or `sqlite://` scheme/prefix", *bundle.DatabaseURL)

‎cmd/river/rivercli/driver_procurer.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,16 @@ type DriverProcurer interface {
2323
QueryRow(ctx context.Context, sql string, args ...any) riverdriver.Row
2424
}
2525

26+
// DriverProcurerPgxV5 optionally initializes a custom procurer with a pgx/v5 pool.
2627
type DriverProcurerPgxV5 interface {
2728
InitPgxV5(pool *pgxpool.Pool)
2829
}
2930

31+
// DriverProcurerSQLite optionally initializes a custom procurer with a SQLite pool.
32+
type DriverProcurerSQLite interface {
33+
InitSQLite(pool *sql.DB)
34+
}
35+
3036
// BenchmarkerInterface is an interface to a Benchmarker. Its reason for
3137
// existence is to wrap a benchmarker to strip it of its generic parameter,
3238
// letting us pass it around without having to know the transaction type.

‎cmd/river/rivercli/river_cli_test.go‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"bytes"
55
"cmp"
66
"context"
7+
"database/sql"
78
"fmt"
89
"maps"
910
"net/url"
@@ -28,6 +29,7 @@ type DriverProcurerStub struct {
2829
getBenchmarkerStub func(config *riverbench.Config) BenchmarkerInterface
2930
getMigratorStub func(config *rivermigrate.Config) (MigratorInterface, error)
3031
initPgxV5Stub func(pool *pgxpool.Pool)
32+
initSQLiteStub func(pool *sql.DB)
3133
queryRowStub func(ctx context.Context, sql string, args ...any) riverdriver.Row
3234
}
3335

@@ -55,6 +57,14 @@ func (p *DriverProcurerStub) InitPgxV5(pool *pgxpool.Pool) {
5557
p.initPgxV5Stub(pool)
5658
}
5759

60+
func (p *DriverProcurerStub) InitSQLite(pool *sql.DB) {
61+
if p.initSQLiteStub == nil {
62+
panic("InitSQLite is not stubbed")
63+
}
64+
65+
p.initSQLiteStub(pool)
66+
}
67+
5868
func (p *DriverProcurerStub) QueryRow(ctx context.Context, sql string, args ...any) riverdriver.Row {
5969
if p.queryRowStub == nil {
6070
panic("QueryRow is not stubbed")
@@ -447,6 +457,52 @@ SELECT 'up 1' FROM river_table
447457
`), strings.TrimSpace(out.String()))
448458
}
449459

460+
func TestBaseCommandSetDriverProcurerSQLite(t *testing.T) {
461+
t.Parallel()
462+
463+
getMigratorCalled := false
464+
initSQLiteCalled := false
465+
466+
migratorStub := &MigratorStub{}
467+
migratorStub.allVersionsStub = func() []rivermigrate.Migration { return []rivermigrate.Migration{testMigration01} }
468+
migratorStub.getVersionStub = func(version int) (rivermigrate.Migration, error) {
469+
if version == 1 {
470+
return testMigration01, nil
471+
}
472+
473+
return rivermigrate.Migration{}, fmt.Errorf("unknown version: %d", version)
474+
}
475+
migratorStub.existingVersionsStub = func(ctx context.Context) ([]rivermigrate.Migration, error) { return nil, nil }
476+
477+
cli := NewCLI(&Config{
478+
DriverProcurer: &DriverProcurerStub{
479+
getMigratorStub: func(config *rivermigrate.Config) (MigratorInterface, error) {
480+
getMigratorCalled = true
481+
return migratorStub, nil
482+
},
483+
initSQLiteStub: func(pool *sql.DB) {
484+
initSQLiteCalled = true
485+
},
486+
},
487+
Name: "River",
488+
})
489+
490+
var out bytes.Buffer
491+
cli.SetOut(&out)
492+
493+
cmd := cli.BaseCommandSet()
494+
cmd.SetArgs([]string{"migrate-get", "--up", "--version", "1", "--database-url", "sqlite://"})
495+
require.NoError(t, cmd.Execute())
496+
497+
require.True(t, getMigratorCalled)
498+
require.True(t, initSQLiteCalled)
499+
500+
require.Equal(t, strings.TrimSpace(`
501+
-- River main migration 001 [up]
502+
SELECT 'up 1' FROM river_table
503+
`), strings.TrimSpace(out.String()))
504+
}
505+
450506
func TestMigrateGet(t *testing.T) {
451507
t.Parallel()
452508

0 commit comments

Comments
 (0)