From 2b9a430e59e6b13c0731190e07dac6da72b1a7fe Mon Sep 17 00:00:00 2001 From: Brandur Date: Tue, 11 Aug 2026 14:29:15 -0500 Subject: [PATCH] Support Yugabyte as a Postgres target This one's aimed at #1346, in which it might be possible for us to support Yugabyte as a database target without a hugely inordinate amount of work. Yugabyte is currently targeting compatibility against Postgres 15 [1]. It doesn't support `xmax` which is what #1346 is about, but somewhat surprisingly, we only use `xmax` in one place and don't use any other Postgres 16+ features (as Postgres 15 is still a valid target in the CI matrix). The `xmax` trick to determine whether an upserted row is new or existing is a little outdated anyway because Postgres 18 added the capability to detect an existing row with `OLD.id IS NOT NULL` [2]. Long run, we should switch to that for everything. Shorter term, Postgres 18 is still quite new, so I propose we do something like this: * If on Postgres 18+ (we should be getting Postgres 19 soon), use `OLD.id IS NOT NULL`. * If on Yugabyte, fall back to the same trick we use in SQLite by upserting rows with a unique nonce and checking whether the nonce was the one we inserted or not. * Otherwise, use the existing approach with `xmax`. We do have to check which database we're on, but only once, after which we can cache that information forever, so it shouldn't have any impact on performance. Fixes #1346. [1] https://docs.yugabyte.com/stable/faq/compatibility/#what-is-the-extent-of-compatibility-with-postgresql [2] https://www.crunchydata.com/blog/postgres-18-old-and-new-in-the-returning-clause --- client.go | 6 +- client_test.go | 20 +++ internal/rivercommon/river_common.go | 5 - riverdriver/river_driver_interface.go | 9 ++ .../internal/dbsqlc/pg_misc.sql.go | 18 +++ .../internal/dbsqlc/river_job.sql.go | 4 +- .../river_database_sql_driver.go | 72 +++++++++- .../river_database_sql_driver_test.go | 61 +++++++++ .../yugabyte_compatibility_test.go | 57 ++++++++ riverdriver/riverdrivertest/job_insert.go | 2 +- riverdriver/riverdrivertest/job_read.go | 6 +- .../riverdrivertest/riverdrivertest.go | 20 +++ .../riverpgxv5/internal/dbsqlc/pg_misc.sql | 5 + .../riverpgxv5/internal/dbsqlc/pg_misc.sql.go | 18 +++ .../riverpgxv5/internal/dbsqlc/river_job.sql | 4 +- .../internal/dbsqlc/river_job.sql.go | 4 +- riverdriver/riverpgxv5/river_pgx_v5_driver.go | 70 +++++++++- .../riverpgxv5/river_pgx_v5_driver_test.go | 55 ++++++++ .../riverpgxv5/yugabyte_compatibility_test.go | 57 ++++++++ riverdriver/riversqlite/go.mod | 8 +- riverdriver/riversqlite/go.sum | 11 -- .../riversqlite/river_sqlite_driver.go | 14 +- riverdriver/unique_insert.go | 112 ++++++++++++++++ riverdriver/unique_insert_test.go | 126 ++++++++++++++++++ 24 files changed, 719 insertions(+), 45 deletions(-) create mode 100644 riverdriver/riverdatabasesql/yugabyte_compatibility_test.go create mode 100644 riverdriver/riverpgxv5/yugabyte_compatibility_test.go create mode 100644 riverdriver/unique_insert.go create mode 100644 riverdriver/unique_insert_test.go diff --git a/client.go b/client.go index bf07bfc0f..2af5092b4 100644 --- a/client.go +++ b/client.go @@ -1110,9 +1110,13 @@ func (c *Client[TTx]) Start(ctx context.Context) error { // available, the client appears to have started even though it's completely // non-functional. Here we try to make an initial assessment of health and // return quickly in case of an apparent problem. - if err := c.driver.GetExecutor().Exec(fetchCtx, "SELECT 1"); err != nil { + executor := c.driver.GetExecutor() + if err := executor.Ping(fetchCtx); err != nil { return fmt.Errorf("error making initial connection to database: %w", err) } + if err := executor.InitDriver(fetchCtx); err != nil { + return fmt.Errorf("error initializing driver: %w", err) + } // Each time we start, we need a fresh completer subscribe channel to // send job completion events on, because the completer will close it diff --git a/client_test.go b/client_test.go index bb7ba8150..0c6ffd23f 100644 --- a/client_test.go +++ b/client_test.go @@ -8316,6 +8316,26 @@ func Test_Client_Start_Error(t *testing.T) { require.Equal(t, pgerrcode.InvalidCatalogName, pgErr.Code) }) + t.Run("DatabaseErrorAfterSuccessfulStart", func(t *testing.T) { + t.Parallel() + + dbPool := riversharedtest.DBPoolClone(ctx, t) + driver := NewDriverPollOnly(dbPool) + schema := riverdbtest.TestSchema(ctx, t, driver, nil) + + client, err := NewClient(driver, newTestConfig(t, schema)) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, client.Stop(ctx)) }) + + require.NoError(t, client.Start(ctx)) + require.NoError(t, client.Stop(ctx)) + + dbPool.Close() + + err = client.Start(ctx) + require.ErrorIs(t, err, riverdriver.ErrClosedPool) + }) + t.Run("CanRestartAfterFailure", func(t *testing.T) { t.Parallel() diff --git a/internal/rivercommon/river_common.go b/internal/rivercommon/river_common.go index 4f769b391..389efdc8a 100644 --- a/internal/rivercommon/river_common.go +++ b/internal/rivercommon/river_common.go @@ -45,11 +45,6 @@ const ( // MetadataKeyRescueCount records how many times the job has been rescued. MetadataKeyRescueCount = "river:rescue_count" - - // MetadataKeyUniqueNonce is a special metadata key used by the SQLite driver to - // determine whether an upsert is was skipped or not because the `(xmax != 0)` - // trick we use in Postgres doesn't work in SQLite. - MetadataKeyUniqueNonce = "river:unique_nonce" ) type ContextKeyClient struct{} diff --git a/riverdriver/river_driver_interface.go b/riverdriver/river_driver_interface.go index d4b87f315..54a1ea1f5 100644 --- a/riverdriver/river_driver_interface.go +++ b/riverdriver/river_driver_interface.go @@ -224,6 +224,11 @@ type Executor interface { IndexReindex(ctx context.Context, params *IndexReindexParams) error IndexReindexArtifacts(ctx context.Context, params *IndexReindexArtifactsParams) ([]string, error) + // InitDriver initializes driver-specific state using information read from + // the database. Implementations must be safe to call concurrently and + // repeatedly, and should cache successfully initialized state. + InitDriver(ctx context.Context) error + JobCancel(ctx context.Context, params *JobCancelParams) (*rivertype.JobRow, error) JobCountByAllStates(ctx context.Context, params *JobCountByAllStatesParams) (map[rivertype.JobState]int, error) JobCountByQueueAndState(ctx context.Context, params *JobCountByQueueAndStateParams) ([]*JobCountByQueueAndStateResult, error) @@ -289,6 +294,10 @@ type Executor interface { NotificationDeleteBefore(ctx context.Context, params *NotificationDeleteBeforeParams) (int, error) NotifyMany(ctx context.Context, params *NotifyManyParams) error + + // Ping checks that the database is reachable. + Ping(ctx context.Context) error + PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) QueueCreateOrSetUpdatedAt(ctx context.Context, params *QueueCreateOrSetUpdatedAtParams) (*rivertype.Queue, error) diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go index e2542e763..b4be68757 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go @@ -21,6 +21,24 @@ func (q *Queries) PGAdvisoryXactLock(ctx context.Context, db DBTX, key int64) er return err } +const pGGetProductAndVersion = `-- name: PGGetProductAndVersion :one +SELECT + version()::text AS product, + current_setting('server_version_num')::int AS version_num +` + +type PGGetProductAndVersionRow struct { + Product string + VersionNum int32 +} + +func (q *Queries) PGGetProductAndVersion(ctx context.Context, db DBTX) (*PGGetProductAndVersionRow, error) { + row := db.QueryRowContext(ctx, pGGetProductAndVersion) + var i PGGetProductAndVersionRow + err := row.Scan(&i.Product, &i.VersionNum) + return &i, err +} + const pGNotifyMany = `-- name: PGNotifyMany :exec WITH topic_to_notify AS ( SELECT diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go index a197f7bb5..e7361009e 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go @@ -718,7 +718,9 @@ ON CONFLICT (unique_key) AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state) -- Something needs to be updated for a row to be returned on a conflict. DO UPDATE SET kind = EXCLUDED.kind -RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states, (xmax != 0) AS unique_skipped_as_duplicate +RETURNING + river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states, + /* TEMPLATE_BEGIN: unique_skipped_as_duplicate */ (xmax != 0) /* TEMPLATE_END */ AS unique_skipped_as_duplicate ` type JobInsertFastManyParams struct { diff --git a/riverdriver/riverdatabasesql/river_database_sql_driver.go b/riverdriver/riverdatabasesql/river_database_sql_driver.go index b672c77d0..4902b1a1c 100644 --- a/riverdriver/riverdatabasesql/river_database_sql_driver.go +++ b/riverdriver/riverdatabasesql/river_database_sql_driver.go @@ -16,6 +16,7 @@ import ( "io/fs" "math" "strings" + "sync/atomic" "time" "github.com/jackc/pgx/v5/pgxpool" @@ -28,6 +29,7 @@ import ( "github.com/riverqueue/river/rivershared/uniquestates" "github.com/riverqueue/river/rivershared/util/dbutil" "github.com/riverqueue/river/rivershared/util/ptrutil" + "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/savepointutil" "github.com/riverqueue/river/rivershared/util/sliceutil" "github.com/riverqueue/river/rivertype" @@ -38,9 +40,10 @@ var migrationFS embed.FS // Driver is an implementation of riverdriver.Driver for database/sql. type Driver struct { - dbPool *sql.DB - listenerDriver *riverpgxv5.Driver - replacer sqlctemplate.Replacer + dbPool *sql.DB + listenerDriver *riverpgxv5.Driver + replacer sqlctemplate.Replacer + uniqueInsertMode atomic.Uint32 } // New returns a new database/sql River driver for use with River. @@ -247,6 +250,11 @@ func (e *Executor) IndexesExist(ctx context.Context, params *riverdriver.Indexes return exists, nil } +func (e *Executor) InitDriver(ctx context.Context) error { + _, err := e.uniqueInsertMode(ctx) + return err +} + func (e *Executor) JobCancel(ctx context.Context, params *riverdriver.JobCancelParams) (*rivertype.JobRow, error) { cancelledAt, err := params.CancelAttemptedAt.MarshalJSON() if err != nil { @@ -402,6 +410,16 @@ func (e *Executor) JobGetStuck(ctx context.Context, params *riverdriver.JobGetSt } func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) { + uniqueInsertMode, err := e.uniqueInsertMode(ctx) + if err != nil { + return nil, err + } + + var uniqueNonce string + if uniqueInsertMode == riverdriver.UniqueInsertModeMetadataNonce { + uniqueNonce = randutil.Hex(8) + } + insertJobsParams := &dbsqlc.JobInsertFastManyParams{ ID: make([]int64, len(params.Jobs)), Args: make([]string, len(params.Jobs)), @@ -442,7 +460,16 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo insertJobsParams.CreatedAt[i] = createdAt insertJobsParams.Kind[i] = params.Kind insertJobsParams.MaxAttempts[i] = int16(min(params.MaxAttempts, math.MaxInt16)) //nolint:gosec - insertJobsParams.Metadata[i] = cmp.Or(string(params.Metadata), "{}") + metadata := []byte(cmp.Or(string(params.Metadata), "{}")) + if uniqueNonce != "" { + var err error + metadata, err = riverdriver.UniqueInsertMetadataWithNonce(metadata, uniqueNonce) + if err != nil { + return nil, err + } + } + + insertJobsParams.Metadata[i] = string(metadata) insertJobsParams.Priority[i] = int16(min(params.Priority, math.MaxInt16)) //nolint:gosec insertJobsParams.Queue[i] = params.Queue insertJobsParams.ScheduledAt[i] = scheduledAt @@ -452,6 +479,9 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo insertJobsParams.UniqueStates[i] = int32(params.UniqueStates) } + ctx = sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{ + "unique_skipped_as_duplicate": {Value: uniqueInsertMode.SQL(), Stable: true}, + }, nil) items, err := dbsqlc.New().JobInsertFastMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, insertJobsParams) if err != nil { return nil, interpretError(err) @@ -462,7 +492,13 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo if err != nil { return nil, err } - return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: row.UniqueSkippedAsDuplicate}, nil + + uniqueSkippedAsDuplicate := row.UniqueSkippedAsDuplicate + if uniqueInsertMode == riverdriver.UniqueInsertModeMetadataNonce { + uniqueSkippedAsDuplicate = riverdriver.UniqueInsertMetadataIsDuplicate(job.Metadata, uniqueNonce) + } + + return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: uniqueSkippedAsDuplicate}, nil }) } @@ -932,6 +968,10 @@ func (e *Executor) NotifyMany(ctx context.Context, params *riverdriver.NotifyMan }) } +func (e *Executor) Ping(ctx context.Context) error { + return e.Exec(ctx, "SELECT 1") +} + func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) { err := dbsqlc.New().PGAdvisoryXactLock(ctx, e.dbtx, key) return &struct{}{}, interpretError(err) @@ -1090,6 +1130,28 @@ func (e *Executor) TableTruncate(ctx context.Context, params *riverdriver.TableT return interpretError(err) } +func (e *Executor) uniqueInsertMode(ctx context.Context) (riverdriver.UniqueInsertMode, error) { + if e.driver != nil { + if mode := riverdriver.UniqueInsertMode(e.driver.uniqueInsertMode.Load()); mode != riverdriver.UniqueInsertModeUnknown { + return mode, nil + } + } + + productAndVersion, err := dbsqlc.New().PGGetProductAndVersion(ctx, e.dbtx) + if err != nil { + return riverdriver.UniqueInsertModeUnknown, interpretError(err) + } + + mode := riverdriver.UniqueInsertModeFromProductAndVersion(productAndVersion.Product, productAndVersion.VersionNum) + if e.driver != nil { + // Concurrent callers may both detect, but the first successful result + // becomes the driver's cached mode. + e.driver.uniqueInsertMode.CompareAndSwap(uint32(riverdriver.UniqueInsertModeUnknown), uint32(mode)) + mode = riverdriver.UniqueInsertMode(e.driver.uniqueInsertMode.Load()) + } + return mode, nil +} + type ExecutorTx struct { Executor diff --git a/riverdriver/riverdatabasesql/river_database_sql_driver_test.go b/riverdriver/riverdatabasesql/river_database_sql_driver_test.go index a452dd017..a81d9fe50 100644 --- a/riverdriver/riverdatabasesql/river_database_sql_driver_test.go +++ b/riverdriver/riverdatabasesql/river_database_sql_driver_test.go @@ -5,18 +5,34 @@ import ( "database/sql" "errors" "testing" + "time" "github.com/jackc/pgx/v5/pgxpool" + _ "github.com/jackc/pgx/v5/stdlib" "github.com/stretchr/testify/require" "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/riversharedtest" "github.com/riverqueue/river/rivershared/sqlctemplate" + "github.com/riverqueue/river/rivershared/testsignal" + "github.com/riverqueue/river/rivershared/util/urlutil" "github.com/riverqueue/river/rivertype" ) // Verify interface compliance. var _ riverdriver.Driver[*sql.Tx] = New(nil) +type executorInitDriverTestDBTX struct { + *sql.DB + + QueryRowStarted testsignal.TestSignal[struct{}] +} + +func (d *executorInitDriverTestDBTX) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row { + d.QueryRowStarted.Signal(struct{}{}) + return d.DB.QueryRowContext(ctx, query, args...) +} + func TestNew(t *testing.T) { t.Parallel() @@ -48,6 +64,51 @@ func TestNew(t *testing.T) { }) } +func TestExecutor_InitDriverDoesNotBlockTransaction(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dbPool, err := sql.Open("pgx", urlutil.DatabaseSQLCompatibleURL(riversharedtest.TestDatabaseURL())) + require.NoError(t, err) + dbPool.SetMaxOpenConns(1) + t.Cleanup(func() { require.NoError(t, dbPool.Close()) }) + + driver := New(dbPool) + tx, err := dbPool.BeginTx(ctx, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = tx.Rollback() }) + + poolDBTX := &executorInitDriverTestDBTX{DB: dbPool} + poolDBTX.QueryRowStarted.Init(t) + poolExecutor := &Executor{ + dbPool: dbPool, + dbtx: templateReplaceWrapper{dbtx: poolDBTX, replacer: &driver.replacer}, + driver: driver, + } + + initCtx, initCancel := context.WithTimeout(ctx, 10*time.Second) + t.Cleanup(initCancel) + + var poolInitFinished testsignal.TestSignal[error] + poolInitFinished.Init(t) + go func() { poolInitFinished.Signal(poolExecutor.InitDriver(initCtx)) }() + poolDBTX.QueryRowStarted.WaitOrTimeout() + + var txInitFinished testsignal.TestSignal[error] + txInitFinished.Init(t) + go func() { txInitFinished.Signal(driver.UnwrapExecutor(tx).InitDriver(ctx)) }() + + select { + case err := <-txInitFinished.WaitC(): + require.NoError(t, err) + case <-time.After(2 * time.Second): + require.FailNow(t, "transactional driver initialization blocked behind pool initialization") + } + + require.NoError(t, tx.Rollback()) + require.NoError(t, poolInitFinished.WaitOrTimeout()) +} + func TestNewWithPgxListener(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverdatabasesql/yugabyte_compatibility_test.go b/riverdriver/riverdatabasesql/yugabyte_compatibility_test.go new file mode 100644 index 000000000..a4d9fcf51 --- /dev/null +++ b/riverdriver/riverdatabasesql/yugabyte_compatibility_test.go @@ -0,0 +1,57 @@ +package riverdatabasesql + +import ( + "fmt" + "io/fs" + "os" + "regexp" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestYugabyteCompatibility(t *testing.T) { + t.Parallel() + + // YugabyteDB doesn't expose PostgreSQL's transaction-related system columns: + // https://docs.yugabyte.com/stable/yugabyte-voyager/known-issues/postgresql/#system-columns-is-not-yet-supported + // The unique insert query may use xmax only because the entire expression is + // replaced when the driver detects YugabyteDB. + var ( + uniqueInsertModeTemplateRE = regexp.MustCompile(`(?s)/\*\s*TEMPLATE_BEGIN: unique_skipped_as_duplicate\s*\*/.*?/\*\s*TEMPLATE_END\s*\*/`) + unsupportedSystemColumnRE = regexp.MustCompile(`(?i)\b(?:cmax|cmin|ctid|xmax|xmin)\b`) + ) + + sourceRoot, err := os.OpenRoot(".") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, sourceRoot.Close()) }) + + var violations []string + err = fs.WalkDir(sourceRoot.FS(), ".", func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + if entry.IsDir() || (!strings.HasSuffix(path, ".go") && !strings.HasSuffix(path, ".sql")) || strings.HasSuffix(path, "_test.go") { + return nil + } + + contents, err := sourceRoot.ReadFile(path) + if err != nil { + return err + } + contents = uniqueInsertModeTemplateRE.ReplaceAll(contents, nil) + + for lineNum, line := range strings.Split(string(contents), "\n") { + for _, column := range unsupportedSystemColumnRE.FindAllString(line, -1) { + violations = append(violations, fmt.Sprintf("%s:%d: %s", path, lineNum+1, column)) + } + } + return nil + }) + require.NoError(t, err) + + require.Empty(t, violations, + "YugabyteDB-incompatible PostgreSQL system columns must only appear inside SQL templates that replace them for YugabyteDB", + ) +} diff --git a/riverdriver/riverdrivertest/job_insert.go b/riverdriver/riverdrivertest/job_insert.go index 6b625a0e5..9bc8cf64a 100644 --- a/riverdriver/riverdrivertest/job_insert.go +++ b/riverdriver/riverdrivertest/job_insert.go @@ -88,7 +88,7 @@ func exerciseJobInsert[TTx any](ctx context.Context, t *testing.T, // SQLite needs to set a special metadata key to be able to // check for duplicates. Remove this for purposes of comparing // inserted metadata. - job.Metadata, err = sjson.DeleteBytes(job.Metadata, rivercommon.MetadataKeyUniqueNonce) + job.Metadata, err = sjson.DeleteBytes(job.Metadata, riverdriver.UniqueInsertMetadataKey) require.NoError(t, err) require.Equal(t, idStart+int64(i), job.ID) diff --git a/riverdriver/riverdrivertest/job_read.go b/riverdriver/riverdrivertest/job_read.go index f8fe44182..eaa090827 100644 --- a/riverdriver/riverdrivertest/job_read.go +++ b/riverdriver/riverdrivertest/job_read.go @@ -55,7 +55,7 @@ func exerciseJobRead[TTx any](ctx context.Context, t *testing.T, executorWithTx for _, state := range rivertype.JobStates() { require.Contains(t, countsByState, state) - switch state { //nolint:exhaustive + switch state { case rivertype.JobStateAvailable: require.Equal(t, 2, countsByState[state]) case rivertype.JobStateCancelled: @@ -64,8 +64,10 @@ func exerciseJobRead[TTx any](ctx context.Context, t *testing.T, executorWithTx require.Equal(t, 1, countsByState[state]) case rivertype.JobStateDiscarded: require.Equal(t, 1, countsByState[state]) - default: + case rivertype.JobStatePending, rivertype.JobStateRetryable, rivertype.JobStateRunning, rivertype.JobStateScheduled: require.Equal(t, 0, countsByState[state]) + default: + require.FailNow(t, "unknown job state", state) } } }) diff --git a/riverdriver/riverdrivertest/riverdrivertest.go b/riverdriver/riverdrivertest/riverdrivertest.go index aa66762c7..41c50b1a7 100644 --- a/riverdriver/riverdrivertest/riverdrivertest.go +++ b/riverdriver/riverdrivertest/riverdrivertest.go @@ -54,6 +54,26 @@ func exerciseDriverPool[TTx any](ctx context.Context, t *testing.T, ) { t.Helper() + t.Run("InitDriver", func(t *testing.T) { + t.Parallel() + + exec, _ := executorWithTx(ctx, t) + require.NoError(t, exec.InitDriver(ctx)) + require.NoError(t, exec.InitDriver(ctx)) + }) + + t.Run("Ping", func(t *testing.T) { + t.Parallel() + + exec, _ := executorWithTx(ctx, t) + require.NoError(t, exec.InitDriver(ctx)) + require.NoError(t, exec.Ping(ctx)) + + cancelledCtx, cancel := context.WithCancel(ctx) + cancel() + require.ErrorIs(t, exec.Ping(cancelledCtx), context.Canceled) + }) + t.Run("PoolIsSet", func(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql index 19a7b99f6..02f510b2f 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql @@ -1,6 +1,11 @@ -- name: PGAdvisoryXactLock :exec SELECT pg_advisory_xact_lock(@key); +-- name: PGGetProductAndVersion :one +SELECT + version()::text AS product, + current_setting('server_version_num')::int AS version_num; + -- name: PGNotifyMany :exec WITH topic_to_notify AS ( SELECT diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go index 9215c089b..e1ef0917b 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql.go @@ -20,6 +20,24 @@ func (q *Queries) PGAdvisoryXactLock(ctx context.Context, db DBTX, key int64) er return err } +const pGGetProductAndVersion = `-- name: PGGetProductAndVersion :one +SELECT + version()::text AS product, + current_setting('server_version_num')::int AS version_num +` + +type PGGetProductAndVersionRow struct { + Product string + VersionNum int32 +} + +func (q *Queries) PGGetProductAndVersion(ctx context.Context, db DBTX) (*PGGetProductAndVersionRow, error) { + row := db.QueryRow(ctx, pGGetProductAndVersion) + var i PGGetProductAndVersionRow + err := row.Scan(&i.Product, &i.VersionNum) + return &i, err +} + const pGNotifyMany = `-- name: PGNotifyMany :exec WITH topic_to_notify AS ( SELECT diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql index 40098ed07..ca953ae57 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql @@ -318,7 +318,9 @@ ON CONFLICT (unique_key) AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state) -- Something needs to be updated for a row to be returned on a conflict. DO UPDATE SET kind = EXCLUDED.kind -RETURNING sqlc.embed(river_job), (xmax != 0) AS unique_skipped_as_duplicate; +RETURNING + sqlc.embed(river_job), + /* TEMPLATE_BEGIN: unique_skipped_as_duplicate */ (xmax != 0) /* TEMPLATE_END */ AS unique_skipped_as_duplicate; -- name: JobInsertFastManyNoReturning :execrows INSERT INTO /* TEMPLATE: schema */river_job( diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go index 7eb082c52..0043759b1 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go @@ -697,7 +697,9 @@ ON CONFLICT (unique_key) AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state) -- Something needs to be updated for a row to be returned on a conflict. DO UPDATE SET kind = EXCLUDED.kind -RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states, (xmax != 0) AS unique_skipped_as_duplicate +RETURNING + river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states, + /* TEMPLATE_BEGIN: unique_skipped_as_duplicate */ (xmax != 0) /* TEMPLATE_END */ AS unique_skipped_as_duplicate ` type JobInsertFastManyParams struct { diff --git a/riverdriver/riverpgxv5/river_pgx_v5_driver.go b/riverdriver/riverpgxv5/river_pgx_v5_driver.go index 34954fac0..51eaaccdf 100644 --- a/riverdriver/riverpgxv5/river_pgx_v5_driver.go +++ b/riverdriver/riverpgxv5/river_pgx_v5_driver.go @@ -16,6 +16,7 @@ import ( "math" "strings" "sync" + "sync/atomic" "time" "github.com/jackc/pgx/v5" @@ -30,6 +31,7 @@ import ( "github.com/riverqueue/river/rivershared/uniquestates" "github.com/riverqueue/river/rivershared/util/dbutil" "github.com/riverqueue/river/rivershared/util/ptrutil" + "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/sliceutil" "github.com/riverqueue/river/rivertype" ) @@ -39,8 +41,9 @@ var migrationFS embed.FS // Driver is an implementation of riverdriver.Driver for Pgx v5. type Driver struct { - dbPool *pgxpool.Pool - replacer sqlctemplate.Replacer + dbPool *pgxpool.Pool + replacer sqlctemplate.Replacer + uniqueInsertMode atomic.Uint32 } // New returns a new Pgx v5 River driver for use with River. @@ -216,6 +219,11 @@ func (e *Executor) IndexesExist(ctx context.Context, params *riverdriver.Indexes return exists, nil } +func (e *Executor) InitDriver(ctx context.Context) error { + _, err := e.uniqueInsertMode(ctx) + return err +} + func (e *Executor) JobCancel(ctx context.Context, params *riverdriver.JobCancelParams) (*rivertype.JobRow, error) { cancelledAt, err := params.CancelAttemptedAt.MarshalJSON() if err != nil { @@ -367,6 +375,16 @@ func (e *Executor) JobGetStuck(ctx context.Context, params *riverdriver.JobGetSt } func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) { + uniqueInsertMode, err := e.uniqueInsertMode(ctx) + if err != nil { + return nil, err + } + + var uniqueNonce string + if uniqueInsertMode == riverdriver.UniqueInsertModeMetadataNonce { + uniqueNonce = randutil.Hex(8) + } + insertJobsParams := &dbsqlc.JobInsertFastManyParams{ ID: make([]int64, len(params.Jobs)), Args: make([][]byte, len(params.Jobs)), @@ -408,7 +426,16 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo insertJobsParams.CreatedAt[i] = createdAt insertJobsParams.Kind[i] = params.Kind insertJobsParams.MaxAttempts[i] = int16(min(params.MaxAttempts, math.MaxInt16)) //nolint:gosec - insertJobsParams.Metadata[i] = sliceutil.FirstNonEmpty(params.Metadata, defaultObject) + metadata := sliceutil.FirstNonEmpty(params.Metadata, defaultObject) + if uniqueNonce != "" { + var err error + metadata, err = riverdriver.UniqueInsertMetadataWithNonce(metadata, uniqueNonce) + if err != nil { + return nil, err + } + } + + insertJobsParams.Metadata[i] = metadata insertJobsParams.Priority[i] = int16(min(params.Priority, math.MaxInt16)) //nolint:gosec insertJobsParams.Queue[i] = params.Queue insertJobsParams.ScheduledAt[i] = scheduledAt @@ -418,6 +445,9 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo insertJobsParams.UniqueStates[i] = int32(params.UniqueStates) } + ctx = sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{ + "unique_skipped_as_duplicate": {Value: uniqueInsertMode.SQL(), Stable: true}, + }, nil) items, err := dbsqlc.New().JobInsertFastMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, insertJobsParams) if err != nil { return nil, interpretError(err) @@ -428,7 +458,13 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo if err != nil { return nil, err } - return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: row.UniqueSkippedAsDuplicate}, nil + + uniqueSkippedAsDuplicate := row.UniqueSkippedAsDuplicate + if uniqueInsertMode == riverdriver.UniqueInsertModeMetadataNonce { + uniqueSkippedAsDuplicate = riverdriver.UniqueInsertMetadataIsDuplicate(job.Metadata, uniqueNonce) + } + + return &riverdriver.JobInsertFastResult{Job: job, UniqueSkippedAsDuplicate: uniqueSkippedAsDuplicate}, nil }) } @@ -878,6 +914,10 @@ func (e *Executor) NotifyMany(ctx context.Context, params *riverdriver.NotifyMan }) } +func (e *Executor) Ping(ctx context.Context) error { + return e.Exec(ctx, "SELECT 1") +} + func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) { err := dbsqlc.New().PGAdvisoryXactLock(ctx, e.dbtx, key) return &struct{}{}, interpretError(err) @@ -1036,6 +1076,28 @@ func (e *Executor) TableTruncate(ctx context.Context, params *riverdriver.TableT return interpretError(err) } +func (e *Executor) uniqueInsertMode(ctx context.Context) (riverdriver.UniqueInsertMode, error) { + if e.driver != nil { + if mode := riverdriver.UniqueInsertMode(e.driver.uniqueInsertMode.Load()); mode != riverdriver.UniqueInsertModeUnknown { + return mode, nil + } + } + + productAndVersion, err := dbsqlc.New().PGGetProductAndVersion(ctx, e.dbtx) + if err != nil { + return riverdriver.UniqueInsertModeUnknown, interpretError(err) + } + + mode := riverdriver.UniqueInsertModeFromProductAndVersion(productAndVersion.Product, productAndVersion.VersionNum) + if e.driver != nil { + // Concurrent callers may both detect, but the first successful result + // becomes the driver's cached mode. + e.driver.uniqueInsertMode.CompareAndSwap(uint32(riverdriver.UniqueInsertModeUnknown), uint32(mode)) + mode = riverdriver.UniqueInsertMode(e.driver.uniqueInsertMode.Load()) + } + return mode, nil +} + type ExecutorTx struct { Executor diff --git a/riverdriver/riverpgxv5/river_pgx_v5_driver_test.go b/riverdriver/riverpgxv5/river_pgx_v5_driver_test.go index ae366d5a7..34d20da5a 100644 --- a/riverdriver/riverpgxv5/river_pgx_v5_driver_test.go +++ b/riverdriver/riverpgxv5/river_pgx_v5_driver_test.go @@ -18,12 +18,24 @@ import ( "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/sqlctemplate" + "github.com/riverqueue/river/rivershared/testsignal" "github.com/riverqueue/river/rivertype" ) // Verify interface compliance. var _ riverdriver.Driver[pgx.Tx] = New(nil) +type executorInitDriverTestDBTX struct { + *pgxpool.Pool + + QueryRowStarted testsignal.TestSignal[struct{}] +} + +func (d *executorInitDriverTestDBTX) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { + d.QueryRowStarted.Signal(struct{}{}) + return d.Pool.QueryRow(ctx, sql, args...) +} + func TestNew(t *testing.T) { t.Parallel() @@ -43,6 +55,49 @@ func TestNew(t *testing.T) { }) } +func TestExecutor_InitDriverDoesNotBlockTransaction(t *testing.T) { + t.Parallel() + + ctx := context.Background() + config := testPoolConfig() + config.MaxConns = 1 + dbPool := testPool(ctx, t, config) + driver := New(dbPool) + + tx, err := dbPool.Begin(ctx) + require.NoError(t, err) + t.Cleanup(func() { _ = tx.Rollback(ctx) }) + + poolDBTX := &executorInitDriverTestDBTX{Pool: dbPool} + poolDBTX.QueryRowStarted.Init(t) + poolExecutor := &Executor{ + dbtx: templateReplaceWrapper{dbtx: poolDBTX, replacer: &driver.replacer}, + driver: driver, + } + + initCtx, initCancel := context.WithTimeout(ctx, 10*time.Second) + t.Cleanup(initCancel) + + var poolInitFinished testsignal.TestSignal[error] + poolInitFinished.Init(t) + go func() { poolInitFinished.Signal(poolExecutor.InitDriver(initCtx)) }() + poolDBTX.QueryRowStarted.WaitOrTimeout() + + var txInitFinished testsignal.TestSignal[error] + txInitFinished.Init(t) + go func() { txInitFinished.Signal(driver.UnwrapExecutor(tx).InitDriver(ctx)) }() + + select { + case err := <-txInitFinished.WaitC(): + require.NoError(t, err) + case <-time.After(2 * time.Second): + require.FailNow(t, "transactional driver initialization blocked behind pool initialization") + } + + require.NoError(t, tx.Rollback(ctx)) + require.NoError(t, poolInitFinished.WaitOrTimeout()) +} + func TestListener_Close(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverpgxv5/yugabyte_compatibility_test.go b/riverdriver/riverpgxv5/yugabyte_compatibility_test.go new file mode 100644 index 000000000..57aafa691 --- /dev/null +++ b/riverdriver/riverpgxv5/yugabyte_compatibility_test.go @@ -0,0 +1,57 @@ +package riverpgxv5 + +import ( + "fmt" + "io/fs" + "os" + "regexp" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestYugabyteCompatibility(t *testing.T) { + t.Parallel() + + // YugabyteDB doesn't expose PostgreSQL's transaction-related system columns: + // https://docs.yugabyte.com/stable/yugabyte-voyager/known-issues/postgresql/#system-columns-is-not-yet-supported + // The unique insert query may use xmax only because the entire expression is + // replaced when the driver detects YugabyteDB. + var ( + uniqueInsertModeTemplateRE = regexp.MustCompile(`(?s)/\*\s*TEMPLATE_BEGIN: unique_skipped_as_duplicate\s*\*/.*?/\*\s*TEMPLATE_END\s*\*/`) + unsupportedSystemColumnRE = regexp.MustCompile(`(?i)\b(?:cmax|cmin|ctid|xmax|xmin)\b`) + ) + + sourceRoot, err := os.OpenRoot(".") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, sourceRoot.Close()) }) + + var violations []string + err = fs.WalkDir(sourceRoot.FS(), ".", func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + if entry.IsDir() || (!strings.HasSuffix(path, ".go") && !strings.HasSuffix(path, ".sql")) || strings.HasSuffix(path, "_test.go") { + return nil + } + + contents, err := sourceRoot.ReadFile(path) + if err != nil { + return err + } + contents = uniqueInsertModeTemplateRE.ReplaceAll(contents, nil) + + for lineNum, line := range strings.Split(string(contents), "\n") { + for _, column := range unsupportedSystemColumnRE.FindAllString(line, -1) { + violations = append(violations, fmt.Sprintf("%s:%d: %s", path, lineNum+1, column)) + } + } + return nil + }) + require.NoError(t, err) + + require.Empty(t, violations, + "YugabyteDB-incompatible PostgreSQL system columns must only appear inside SQL templates that replace them for YugabyteDB", + ) +} diff --git a/riverdriver/riversqlite/go.mod b/riverdriver/riversqlite/go.mod index 7af5377bb..23be1c865 100644 --- a/riverdriver/riversqlite/go.mod +++ b/riverdriver/riversqlite/go.mod @@ -10,12 +10,6 @@ require ( github.com/riverqueue/river/rivershared v0.46.0 github.com/riverqueue/river/rivertype v0.46.0 github.com/stretchr/testify v1.12.1 - github.com/tidwall/gjson v1.19.0 - github.com/tidwall/sjson v1.2.5 ) -require ( - github.com/tidwall/match v1.2.0 // indirect - github.com/tidwall/pretty v1.2.1 // indirect - go.yaml.in/yaml/v3 v3.0.5 // indirect -) +require go.yaml.in/yaml/v3 v3.0.5 // indirect diff --git a/riverdriver/riversqlite/go.sum b/riverdriver/riversqlite/go.sum index f66ab1c50..1b7d4e9f5 100644 --- a/riverdriver/riversqlite/go.sum +++ b/riverdriver/riversqlite/go.sum @@ -18,17 +18,6 @@ github.com/riverqueue/river/rivertype v0.44.1 h1:G9UnsBJhlndreMPdjHSoY8wp+OTAsDT github.com/riverqueue/river/rivertype v0.44.1/go.mod h1:D1Ad+EaZiaXbQbJcJcfeicXJMBKno0n6UcfKI5Q7DIQ= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= -github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= -github.com/tidwall/gjson v1.19.0 h1:xwxm7n691Uf3u5OFjzngavjGTh55KX5q/9w9xHW88JU= -github.com/tidwall/gjson v1.19.0/go.mod h1:V37/opeE/JbLUOfH0QTXiNez2l0RUjYUhpT4szFQAfc= -github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= -github.com/tidwall/match v1.2.0 h1:0pt8FlkOwjN2fPt4bIl4BoNxb98gGHN2ObFEDkrfZnM= -github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= -github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= -github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= -github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= -github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= -github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= diff --git a/riverdriver/riversqlite/river_sqlite_driver.go b/riverdriver/riversqlite/river_sqlite_driver.go index c778396d1..524a484c5 100644 --- a/riverdriver/riversqlite/river_sqlite_driver.go +++ b/riverdriver/riversqlite/river_sqlite_driver.go @@ -37,10 +37,6 @@ import ( "sync" "time" - "github.com/tidwall/gjson" - "github.com/tidwall/sjson" - - "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/riverdriver/riversqlite/internal/dbsqlc" "github.com/riverqueue/river/rivershared/sqlctemplate" @@ -293,6 +289,8 @@ func (e *Executor) IndexesExist(ctx context.Context, params *riverdriver.Indexes return exists, nil } +func (e *Executor) InitDriver(context.Context) error { return nil } + func (e *Executor) JobCancel(ctx context.Context, params *riverdriver.JobCancelParams) (*rivertype.JobRow, error) { // Unlike Postgres, this must be carried out in two operations because // SQLite doesn't support CTEs containing `UPDATE`. As long as the job @@ -598,7 +596,7 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo return &riverdriver.JobInsertFastResult{ Job: job, - UniqueSkippedAsDuplicate: gjson.GetBytes(job.Metadata, rivercommon.MetadataKeyUniqueNonce).Str != uniqueNonce, + UniqueSkippedAsDuplicate: riverdriver.UniqueInsertMetadataIsDuplicate(job.Metadata, uniqueNonce), }, nil }) } @@ -1186,6 +1184,10 @@ func (e *Executor) NotifyMany(ctx context.Context, params *riverdriver.NotifyMan return dbsqlc.New().NotificationInsertMany(schemaTemplateParam(ctx, params.Schema), e.dbtx, notifications) } +func (e *Executor) Ping(ctx context.Context) error { + return e.Exec(ctx, "SELECT 1") +} + func (e *Executor) PGAdvisoryXactLock(ctx context.Context, key int64) (*struct{}, error) { return nil, riverdriver.ErrNotImplemented } @@ -1516,7 +1518,7 @@ func sqliteJobInsertFastManyJobsParam(jobs []*riverdriver.JobInsertFastParams, u metadata := sliceutil.FirstNonEmpty(job.Metadata, []byte("{}")) if uniqueNonce != "" { var err error - metadata, err = sjson.SetBytes(metadata, rivercommon.MetadataKeyUniqueNonce, uniqueNonce) + metadata, err = riverdriver.UniqueInsertMetadataWithNonce(metadata, uniqueNonce) if err != nil { return nil, err } diff --git a/riverdriver/unique_insert.go b/riverdriver/unique_insert.go new file mode 100644 index 000000000..191e26da0 --- /dev/null +++ b/riverdriver/unique_insert.go @@ -0,0 +1,112 @@ +package riverdriver + +import ( + "encoding/json" + "fmt" + "strings" +) + +// UniqueInsertMetadataKey is a reserved job metadata key used to detect unique +// insert conflicts on databases that don't expose PostgreSQL system columns. +const UniqueInsertMetadataKey = "river:unique_nonce" + +// UniqueInsertMode is a database-specific strategy for detecting whether a +// unique insert returned a newly inserted job or an existing one. +type UniqueInsertMode uint32 + +const ( + // UniqueInsertModeUnknown indicates that a database's mode hasn't been + // detected yet. + UniqueInsertModeUnknown UniqueInsertMode = iota + + // UniqueInsertModeMetadataNonce detects conflicts by putting a nonce in the + // metadata of the proposed job and checking whether the returned job + // contains it. + UniqueInsertModeMetadataNonce + + // UniqueInsertModeReturningOld uses PostgreSQL 18's OLD row support in + // RETURNING. + UniqueInsertModeReturningOld + + // UniqueInsertModeXmax uses PostgreSQL's xmax system column. + UniqueInsertModeXmax +) + +// SQL returns the SQL expression for the mode. UniqueInsertModeMetadataNonce +// always returns false because duplicate detection is performed in Go instead. +func (m UniqueInsertMode) SQL() string { + switch m { + case UniqueInsertModeMetadataNonce: + return "false" + + case UniqueInsertModeReturningOld: + return "(OLD.id IS NOT NULL)" + + case UniqueInsertModeXmax: + return "(xmax != 0)" + + case UniqueInsertModeUnknown: + panic("unique insert mode has not been detected") + + default: + panic(fmt.Sprintf("invalid unique insert mode: %d", m)) + } +} + +// UniqueInsertMetadataIsDuplicate returns whether metadata lacks the nonce +// from a proposed insert, indicating that an existing row was returned +// instead. +func UniqueInsertMetadataIsDuplicate(metadata []byte, nonce string) bool { + var metadataMap map[string]json.RawMessage + if err := json.Unmarshal(metadata, &metadataMap); err != nil { + return true + } + + var metadataNonce string + if err := json.Unmarshal(metadataMap[UniqueInsertMetadataKey], &metadataNonce); err != nil { + return true + } + return metadataNonce != nonce +} + +// UniqueInsertMetadataWithNonce returns metadata with nonce set under +// UniqueInsertMetadataKey. +func UniqueInsertMetadataWithNonce(metadata []byte, nonce string) ([]byte, error) { + if len(metadata) == 0 { + metadata = []byte("{}") + } + + var metadataMap map[string]json.RawMessage + if err := json.Unmarshal(metadata, &metadataMap); err != nil { + return nil, fmt.Errorf("error unmarshaling job metadata: %w", err) + } + if metadataMap == nil { + metadataMap = make(map[string]json.RawMessage) + } + + nonceJSON, err := json.Marshal(nonce) + if err != nil { + return nil, fmt.Errorf("error marshaling unique insert nonce: %w", err) + } + metadataMap[UniqueInsertMetadataKey] = nonceJSON + + metadata, err = json.Marshal(metadataMap) + if err != nil { + return nil, fmt.Errorf("error marshaling job metadata: %w", err) + } + return metadata, nil +} + +// UniqueInsertModeFromProductAndVersion returns the unique insert mode +// appropriate for a database product and its PostgreSQL-compatible server +// version number. +func UniqueInsertModeFromProductAndVersion(product string, version int32) UniqueInsertMode { + productLower := strings.ToLower(product) + if strings.Contains(productLower, "-yb") || strings.Contains(productLower, "yugabyte") { + return UniqueInsertModeMetadataNonce + } + if version >= 180_000 { + return UniqueInsertModeReturningOld + } + return UniqueInsertModeXmax +} diff --git a/riverdriver/unique_insert_test.go b/riverdriver/unique_insert_test.go new file mode 100644 index 000000000..3c78de985 --- /dev/null +++ b/riverdriver/unique_insert_test.go @@ -0,0 +1,126 @@ +package riverdriver + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestUniqueInsertMetadataIsDuplicate(t *testing.T) { + t.Parallel() + + t.Run("DifferentNonce", func(t *testing.T) { + t.Parallel() + + require.True(t, UniqueInsertMetadataIsDuplicate([]byte(`{"river:unique_nonce":"old"}`), "new")) + }) + + t.Run("InvalidMetadata", func(t *testing.T) { + t.Parallel() + + require.True(t, UniqueInsertMetadataIsDuplicate([]byte(`{`), "nonce")) + }) + + t.Run("MatchingNonce", func(t *testing.T) { + t.Parallel() + + require.False(t, UniqueInsertMetadataIsDuplicate([]byte(`{"river:unique_nonce":"nonce"}`), "nonce")) + }) + + t.Run("MissingNonce", func(t *testing.T) { + t.Parallel() + + require.True(t, UniqueInsertMetadataIsDuplicate([]byte(`{"existing":123}`), "nonce")) + }) +} + +func TestUniqueInsertMetadataWithNonce(t *testing.T) { + t.Parallel() + + t.Run("EmptyMetadata", func(t *testing.T) { + t.Parallel() + + metadata, err := UniqueInsertMetadataWithNonce(nil, "nonce") + require.NoError(t, err) + require.JSONEq(t, `{"river:unique_nonce":"nonce"}`, string(metadata)) + }) + + t.Run("ExistingMetadata", func(t *testing.T) { + t.Parallel() + + metadata, err := UniqueInsertMetadataWithNonce([]byte(`{"existing":123}`), "nonce") + require.NoError(t, err) + require.JSONEq(t, `{"existing":123,"river:unique_nonce":"nonce"}`, string(metadata)) + }) + + t.Run("ExistingNonce", func(t *testing.T) { + t.Parallel() + + metadata, err := UniqueInsertMetadataWithNonce([]byte(`{"river:unique_nonce":"old"}`), "new") + require.NoError(t, err) + require.JSONEq(t, `{"river:unique_nonce":"new"}`, string(metadata)) + }) + + t.Run("InvalidMetadata", func(t *testing.T) { + t.Parallel() + + _, err := UniqueInsertMetadataWithNonce([]byte(`{`), "nonce") + require.ErrorContains(t, err, "error unmarshaling job metadata") + }) +} + +func TestUniqueInsertModeFromProductAndVersion(t *testing.T) { + t.Parallel() + + t.Run("PostgreSQL17", func(t *testing.T) { + t.Parallel() + + require.Equal(t, UniqueInsertModeXmax, UniqueInsertModeFromProductAndVersion("PostgreSQL 17.5", 170_005)) + }) + + t.Run("PostgreSQL18", func(t *testing.T) { + t.Parallel() + + require.Equal(t, UniqueInsertModeReturningOld, UniqueInsertModeFromProductAndVersion("PostgreSQL 18.0", 180_000)) + }) + + t.Run("YugabyteByName", func(t *testing.T) { + t.Parallel() + + require.Equal(t, UniqueInsertModeMetadataNonce, UniqueInsertModeFromProductAndVersion("YugabyteDB", 180_000)) + }) + + t.Run("YugabytePostgreSQLVersion", func(t *testing.T) { + t.Parallel() + + require.Equal(t, UniqueInsertModeMetadataNonce, UniqueInsertModeFromProductAndVersion("PostgreSQL 15.2-YB-2.25.1.0-b0", 150_002)) + }) +} + +func TestUniqueInsertModeSQL(t *testing.T) { + t.Parallel() + + t.Run("MetadataNonce", func(t *testing.T) { + t.Parallel() + + require.Equal(t, "false", UniqueInsertModeMetadataNonce.SQL()) + }) + + t.Run("ReturningOld", func(t *testing.T) { + t.Parallel() + + require.Equal(t, "(OLD.id IS NOT NULL)", UniqueInsertModeReturningOld.SQL()) + }) + + t.Run("Unknown", func(t *testing.T) { + t.Parallel() + + require.PanicsWithValue(t, "unique insert mode has not been detected", func() { UniqueInsertModeUnknown.SQL() }) + }) + + t.Run("Xmax", func(t *testing.T) { + t.Parallel() + + require.Equal(t, "(xmax != 0)", UniqueInsertModeXmax.SQL()) + }) +}