From 6d531666a65f5cee28cf2bb93caae3e7fea34160 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Wed, 26 Aug 2026 20:57:18 -0500 Subject: [PATCH] support mixed drivers in transactional work Worker contexts identify the client that is executing a job, but an application transaction may use a different database abstraction against the same database. Binding the context lookup to the transaction type prevents `JobCompleteTx` and the resumable step helpers from supporting that topology. Use the explicit driver type to unwrap each transaction while obtaining the pilot, schema, and clock from the worker client through a type-independent context view. This preserves the worker's completion semantics while allowing the transaction adapter to differ. Cover completion and resumable step persistence with transaction types that differ from the worker client's transaction type. --- client_context.go | 38 ++++++++++++++++++++++++++++++++++++++ job_complete_tx.go | 30 ++++++++++++++---------------- job_complete_tx_test.go | 32 ++++++++++++++++++++++++++++++++ resumable_step_tx.go | 17 ++++++++--------- resumable_step_tx_test.go | 14 ++++++++++++++ 5 files changed, 106 insertions(+), 25 deletions(-) diff --git a/client_context.go b/client_context.go index 4bbded2e..b4045180 100644 --- a/client_context.go +++ b/client_context.go @@ -3,12 +3,50 @@ package river import ( "context" "errors" + "time" "github.com/riverqueue/river/internal/rivercommon" + "github.com/riverqueue/river/rivershared/riverpilot" ) var errClientNotInContext = errors.New("river: client not found in context, can only be used in a Worker") +type clientContextData struct { + Pilot riverpilot.Pilot + Schema string + Time time.Time +} + +type clientContextProvider interface { + clientContextData() *clientContextData +} + +func clientContextDataFromContext(ctx context.Context) *clientContextData { + client, exists := ctx.Value(rivercommon.ContextKeyClient{}).(clientContextProvider) + if !exists || client == nil { + panic(errClientNotInContext) + } + + data := client.clientContextData() + if data == nil { + panic(errClientNotInContext) + } + + return data +} + +func (c *Client[TTx]) clientContextData() *clientContextData { + if c == nil { + return nil + } + + return &clientContextData{ + Pilot: c.Pilot(), + Schema: c.Schema(), + Time: c.baseService.Time.Now(), + } +} + func withClient[TTx any](ctx context.Context, client *Client[TTx]) context.Context { return context.WithValue(ctx, rivercommon.ContextKeyClient{}, client) } diff --git a/job_complete_tx.go b/job_complete_tx.go index f6e6d00f..625d8d47 100644 --- a/job_complete_tx.go +++ b/job_complete_tx.go @@ -15,9 +15,11 @@ import ( // JobCompleteTx marks the job as completed as part of transaction tx. If tx is // rolled back, the completion will be as well. // -// The function needs to know the type of the River database driver, which is -// the same as the one in use by Client, but the other generic parameters can be -// inferred. An invocation should generally look like: +// The function needs to know the type of the River database driver that is +// compatible with tx. This is usually the same driver used by Client, but may +// differ when a worker and its application transactions use different database +// abstractions. The other generic parameters can be inferred. An invocation +// should generally look like: // // _, err := river.JobCompleteTx[*riverpgxv5.Driver](ctx, tx, job) // if err != nil { @@ -35,31 +37,27 @@ func JobCompleteTx[TDriver riverdriver.Driver[TTx], TTx any, TArgs JobArgs](ctx return nil, errors.New("job must be running") } - client := ClientFromContext[TTx](ctx) - if client == nil { - return nil, errors.New("client not found in context, can only work within a River worker") - } + clientData := clientContextDataFromContext(ctx) - driver := client.Driver() - pilot := client.Pilot() + var driver TDriver // extract metadata updates from context metadataUpdates, hasMetadataUpdates := jobexecutor.MetadataUpdatesFromWorkContext(ctx) hasMetadataUpdates = hasMetadataUpdates && len(metadataUpdates) > 0 var ( metadataUpdatesBytes []byte - err error + marshalErr error ) if hasMetadataUpdates { - metadataUpdatesBytes, err = json.Marshal(metadataUpdates) - if err != nil { - return nil, err + metadataUpdatesBytes, marshalErr = json.Marshal(metadataUpdates) + if marshalErr != nil { + return nil, marshalErr } } execTx := driver.UnwrapExecutor(tx) - params := riverdriver.JobSetStateCompleted(job.ID, client.baseService.Time.Now(), nil) - rows, err := pilot.JobSetStateIfRunningMany(ctx, execTx, &riverdriver.JobSetStateIfRunningManyParams{ + params := riverdriver.JobSetStateCompleted(job.ID, clientData.Time, nil) + rows, err := clientData.Pilot.JobSetStateIfRunningMany(ctx, execTx, &riverdriver.JobSetStateIfRunningManyParams{ ID: []int64{params.ID}, Attempt: []*int{params.Attempt}, ErrData: [][]byte{params.ErrData}, @@ -67,7 +65,7 @@ func JobCompleteTx[TDriver riverdriver.Driver[TTx], TTx any, TArgs JobArgs](ctx MetadataDoMerge: []bool{hasMetadataUpdates}, MetadataUpdates: [][]byte{metadataUpdatesBytes}, ScheduledAt: []*time.Time{params.ScheduledAt}, - Schema: client.config.Schema, + Schema: clientData.Schema, State: []rivertype.JobState{params.State}, }) if err != nil { diff --git a/job_complete_tx_test.go b/job_complete_tx_test.go index c5e61e6e..2d3ed055 100644 --- a/job_complete_tx_test.go +++ b/job_complete_tx_test.go @@ -20,6 +20,24 @@ import ( "github.com/riverqueue/river/rivertype" ) +type wrappedPgxTx struct { + pgx.Tx +} + +type wrappedPgxTxDriver struct { + *riverpgxv5.Driver +} + +func (d *wrappedPgxTxDriver) UnwrapExecutor(tx *wrappedPgxTx) riverdriver.ExecutorTx { + var driver *riverpgxv5.Driver + return driver.UnwrapExecutor(tx.Tx) +} + +func (d *wrappedPgxTxDriver) UnwrapTx(execTx riverdriver.ExecutorTx) *wrappedPgxTx { + var driver *riverpgxv5.Driver + return &wrappedPgxTx{Tx: driver.UnwrapTx(execTx)} +} + func TestJobCompleteTx(t *testing.T) { t.Parallel() @@ -162,4 +180,18 @@ func TestJobCompleteTx(t *testing.T) { require.NoError(t, err) }) }) + + t.Run("UsesTransactionDriverInsteadOfWorkerDriver", func(t *testing.T) { + t.Parallel() + + ctx, bundle := setup(ctx, t) + + job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{ + State: new(rivertype.JobStateRunning), + }) + + completedJob, err := JobCompleteTx[*wrappedPgxTxDriver](ctx, &wrappedPgxTx{Tx: bundle.tx}, &Job[JobArgs]{JobRow: job}) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateCompleted, completedJob.State) + }) } diff --git a/resumable_step_tx.go b/resumable_step_tx.go index fddb8c63..6de35947 100644 --- a/resumable_step_tx.go +++ b/resumable_step_tx.go @@ -26,7 +26,7 @@ import ( // Must be called from within a ResumableStep or ResumableStepCursor callback. // The current step name to persist is read from context. func ResumableSetStepTx[TDriver riverdriver.Driver[TTx], TTx any, TArgs JobArgs](ctx context.Context, tx TTx, job *Job[TArgs]) (*Job[TArgs], error) { - return resumableSetStepTx(ctx, tx, job, nil) + return resumableSetStepTx[TDriver](ctx, tx, job, nil) } // ResumableSetStepCursorTx immediately persists the current resumable step and @@ -48,10 +48,10 @@ func ResumableSetStepCursorTx[TDriver riverdriver.Driver[TTx], TTx any, TArgs Jo return nil, err } - return resumableSetStepTx(ctx, tx, job, json.RawMessage(cursorBytes)) + return resumableSetStepTx[TDriver](ctx, tx, job, json.RawMessage(cursorBytes)) } -func resumableSetStepTx[TTx any, TArgs JobArgs](ctx context.Context, tx TTx, job *Job[TArgs], cursor json.RawMessage) (*Job[TArgs], error) { +func resumableSetStepTx[TDriver riverdriver.Driver[TTx], TTx any, TArgs JobArgs](ctx context.Context, tx TTx, job *Job[TArgs], cursor json.RawMessage) (*Job[TArgs], error) { if job.State != rivertype.JobStateRunning { return nil, errors.New("job must be running") } @@ -66,10 +66,9 @@ func resumableSetStepTx[TTx any, TArgs JobArgs](ctx context.Context, tx TTx, job step := state.StepName - client := ClientFromContext[TTx](ctx) - if client == nil { - return nil, errors.New("client not found in context, can only work within a River worker") - } + clientData := clientContextDataFromContext(ctx) + + var driver TDriver metadataUpdates := map[string]any{ rivercommon.MetadataKeyResumableStep: step, @@ -99,11 +98,11 @@ func resumableSetStepTx[TTx any, TArgs JobArgs](ctx context.Context, tx TTx, job return nil, err } - updatedJob, err := client.Driver().UnwrapExecutor(tx).JobUpdate(ctx, &riverdriver.JobUpdateParams{ + updatedJob, err := driver.UnwrapExecutor(tx).JobUpdate(ctx, &riverdriver.JobUpdateParams{ ID: job.ID, MetadataDoMerge: true, Metadata: metadataUpdatesBytes, - Schema: client.config.Schema, + Schema: clientData.Schema, }) if err != nil { if errors.Is(err, rivertype.ErrNotFound) { diff --git a/resumable_step_tx_test.go b/resumable_step_tx_test.go index da1c3280..b12aa7c1 100644 --- a/resumable_step_tx_test.go +++ b/resumable_step_tx_test.go @@ -165,4 +165,18 @@ func TestResumableSetStepTx(t *testing.T) { require.NoError(t, err) }) }) + + t.Run("UsesTransactionDriverInsteadOfWorkerDriver", func(t *testing.T) { + t.Parallel() + + ctx, bundle := setup(ctx, t, "step1") + + job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{ + State: new(rivertype.JobStateRunning), + }) + + updatedJob, err := ResumableSetStepTx[*wrappedPgxTxDriver](ctx, &wrappedPgxTx{Tx: bundle.tx}, &Job[JobArgs]{JobRow: job}) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRunning, updatedJob.State) + }) }