From a524293424b95c5e1417c295b9e566dbc89a090b 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 --- internal/rivercommon/river_common.go | 5 - .../internal/dbsqlc/pg_misc.sql.go | 18 +++ .../internal/dbsqlc/river_job.sql.go | 4 +- .../river_database_sql_driver.go | 67 +++++++++- riverdriver/riverdrivertest/job_insert.go | 2 +- .../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 | 66 ++++++++- riverdriver/riversqlite/go.mod | 5 - riverdriver/riversqlite/go.sum | 11 -- .../riversqlite/river_sqlite_driver.go | 8 +- riverdriver/unique_insert.go | 112 ++++++++++++++++ riverdriver/unique_insert_test.go | 126 ++++++++++++++++++ 15 files changed, 416 insertions(+), 39 deletions(-) create mode 100644 riverdriver/unique_insert.go create mode 100644 riverdriver/unique_insert_test.go diff --git a/internal/rivercommon/river_common.go b/internal/rivercommon/river_common.go index 4f769b39..389efdc8 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/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/pg_misc.sql.go index e2542e76..b4be6875 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 c72d277a..05f96485 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 a37e9a88..8a5b73c8 100644 --- a/riverdriver/riverdatabasesql/river_database_sql_driver.go +++ b/riverdriver/riverdatabasesql/river_database_sql_driver.go @@ -16,6 +16,8 @@ import ( "io/fs" "math" "strings" + "sync" + "sync/atomic" "time" "github.com/lib/pq" @@ -26,6 +28,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" @@ -36,8 +39,10 @@ var migrationFS embed.FS // Driver is an implementation of riverdriver.Driver for database/sql. type Driver struct { - dbPool *sql.DB - replacer sqlctemplate.Replacer + dbPool *sql.DB + replacer sqlctemplate.Replacer + uniqueInsertMode atomic.Uint32 + uniqueInsertModeInitMu sync.Mutex } // New returns a new database/sql River driver for use with River. @@ -362,6 +367,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)), @@ -402,7 +417,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 @@ -412,6 +436,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) @@ -422,7 +449,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 }) } @@ -1050,6 +1083,32 @@ 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 + } + + e.driver.uniqueInsertModeInitMu.Lock() + defer e.driver.uniqueInsertModeInitMu.Unlock() + + 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 { + e.driver.uniqueInsertMode.Store(uint32(mode)) + } + return mode, nil +} + type ExecutorTx struct { Executor diff --git a/riverdriver/riverdrivertest/job_insert.go b/riverdriver/riverdrivertest/job_insert.go index ee4fa745..9fe6fbad 100644 --- a/riverdriver/riverdrivertest/job_insert.go +++ b/riverdriver/riverdrivertest/job_insert.go @@ -89,7 +89,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/riverpgxv5/internal/dbsqlc/pg_misc.sql b/riverdriver/riverpgxv5/internal/dbsqlc/pg_misc.sql index 19a7b99f..02f510b2 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 9215c089..e1ef0917 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 509e479b..5a88af82 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 a361baac..c09b3021 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 71dcc6d5..105bb0be 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,10 @@ 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 + uniqueInsertModeInitMu sync.Mutex } // New returns a new Pgx v5 River driver for use with River. @@ -366,6 +370,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)), @@ -407,7 +421,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 @@ -417,6 +440,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) @@ -427,7 +453,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 }) } @@ -1035,6 +1067,32 @@ 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 + } + + e.driver.uniqueInsertModeInitMu.Lock() + defer e.driver.uniqueInsertModeInitMu.Unlock() + + 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 { + e.driver.uniqueInsertMode.Store(uint32(mode)) + } + return mode, nil +} + type ExecutorTx struct { Executor diff --git a/riverdriver/riversqlite/go.mod b/riverdriver/riversqlite/go.mod index be7b4859..b66164e0 100644 --- a/riverdriver/riversqlite/go.mod +++ b/riverdriver/riversqlite/go.mod @@ -5,19 +5,14 @@ go 1.25.0 toolchain go1.25.7 require ( - github.com/riverqueue/river v0.43.0 github.com/riverqueue/river/riverdriver v0.43.0 github.com/riverqueue/river/rivershared v0.43.0 github.com/riverqueue/river/rivertype v0.43.0 github.com/stretchr/testify v1.11.1 - github.com/tidwall/gjson v1.19.0 - github.com/tidwall/sjson v1.2.5 ) require ( github.com/davecgh/go-spew v1.1.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect - github.com/tidwall/match v1.2.0 // indirect - github.com/tidwall/pretty v1.2.1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/riverdriver/riversqlite/go.sum b/riverdriver/riversqlite/go.sum index 54fd8d2e..c0424607 100644 --- a/riverdriver/riversqlite/go.sum +++ b/riverdriver/riversqlite/go.sum @@ -28,17 +28,6 @@ github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= -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= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= diff --git a/riverdriver/riversqlite/river_sqlite_driver.go b/riverdriver/riversqlite/river_sqlite_driver.go index df1f4acf..f84d55b8 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" @@ -598,7 +594,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 }) } @@ -1526,7 +1522,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 00000000..191e26da --- /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 00000000..3c78de98 --- /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()) + }) +}