diff --git a/client_pilot_test.go b/client_pilot_test.go index fc9bc160c..d52654054 100644 --- a/client_pilot_test.go +++ b/client_pilot_test.go @@ -74,9 +74,9 @@ func (p *pilotSpy) JobGetAvailable(ctx context.Context, exec riverdriver.Executo return p.StandardPilot.JobGetAvailable(ctx, exec, state, params) } -func (p *pilotSpy) JobInsertMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) { +func (p *pilotSpy) JobInsertMany(ctx context.Context, execTx riverdriver.ExecutorTx, params *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) { p.jobInsertManyCalls.Add(1) - return p.StandardPilot.JobInsertMany(ctx, exec, params) + return p.StandardPilot.JobInsertMany(ctx, execTx, params) } func (p *pilotSpy) JobRetry(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRetryParams) (*rivertype.JobRow, error) { diff --git a/rivershared/riverpilot/pilot.go b/rivershared/riverpilot/pilot.go index 5c14cd86f..25108d29e 100644 --- a/rivershared/riverpilot/pilot.go +++ b/rivershared/riverpilot/pilot.go @@ -38,9 +38,11 @@ type Pilot interface { params *riverdriver.JobGetAvailableParams, ) (*riverdriver.JobGetAvailableResult, error) + // JobInsertMany inserts jobs in the supplied transaction. The caller owns + // its commit or rollback; implementations must keep related writes in it. JobInsertMany( ctx context.Context, - exec riverdriver.Executor, + execTx riverdriver.ExecutorTx, params *riverdriver.JobInsertFastManyParams, ) ([]*riverdriver.JobInsertFastResult, error) diff --git a/rivershared/riverpilot/standard_pilot.go b/rivershared/riverpilot/standard_pilot.go index 12cf4ee2e..2e9b9aec5 100644 --- a/rivershared/riverpilot/standard_pilot.go +++ b/rivershared/riverpilot/standard_pilot.go @@ -37,10 +37,10 @@ func (p *StandardPilot) JobCancel(ctx context.Context, exec riverdriver.Executor func (p *StandardPilot) JobInsertMany( ctx context.Context, - exec riverdriver.Executor, + execTx riverdriver.ExecutorTx, params *riverdriver.JobInsertFastManyParams, ) ([]*riverdriver.JobInsertFastResult, error) { - return exec.JobInsertFastMany(ctx, params) + return execTx.JobInsertFastMany(ctx, params) } func (p *StandardPilot) JobRescueMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRescueManyParams) (*struct{}, error) { diff --git a/rivershared/riverpilot/standard_pilot_test.go b/rivershared/riverpilot/standard_pilot_test.go index fc1c350f9..09e57796f 100644 --- a/rivershared/riverpilot/standard_pilot_test.go +++ b/rivershared/riverpilot/standard_pilot_test.go @@ -20,6 +20,34 @@ func (m *standardPilotExecutorMock) JobGetAvailable(ctx context.Context, params return m.jobGetAvailableFunc(ctx, params) } +type standardPilotExecutorTxMock struct { + riverdriver.ExecutorTx + + beginCalls int + commitCalls int + insertFunc func(context.Context, *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) + rollbackCalls int +} + +func (m *standardPilotExecutorTxMock) Begin(context.Context) (riverdriver.ExecutorTx, error) { + m.beginCalls++ + return m, nil +} + +func (m *standardPilotExecutorTxMock) Commit(context.Context) error { + m.commitCalls++ + return nil +} + +func (m *standardPilotExecutorTxMock) JobInsertFastMany(ctx context.Context, params *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) { + return m.insertFunc(ctx, params) +} + +func (m *standardPilotExecutorTxMock) Rollback(context.Context) error { + m.rollbackCalls++ + return nil +} + func TestStandardPilot_JobGetAvailable(t *testing.T) { t.Parallel() @@ -66,3 +94,47 @@ func TestStandardPilot_JobGetAvailable(t *testing.T) { require.ErrorIs(t, err, parentErr) }) } + +func TestStandardPilot_JobInsertMany(t *testing.T) { + t.Parallel() + + for _, fail := range []bool{true, false} { + name := "Success" + if fail { + name = "Error" + } + t.Run(name, func(t *testing.T) { + t.Parallel() + + ctx := context.Background() + params := &riverdriver.JobInsertFastManyParams{} + result := []*riverdriver.JobInsertFastResult{{}} + insertErr := errors.New("insert failure") + calls := 0 + execTx := &standardPilotExecutorTxMock{ + insertFunc: func(innerCtx context.Context, innerParams *riverdriver.JobInsertFastManyParams) ([]*riverdriver.JobInsertFastResult, error) { + calls++ + require.Equal(t, ctx, innerCtx) + require.Same(t, params, innerParams) + if fail { + return nil, insertErr + } + return result, nil + }, + } + + rows, err := (&StandardPilot{}).JobInsertMany(ctx, execTx, params) + if fail { + require.ErrorIs(t, err, insertErr) + require.Nil(t, rows) + } else { + require.NoError(t, err) + require.Equal(t, result, rows) + } + require.Equal(t, 1, calls) + require.Zero(t, execTx.beginCalls, "the pilot uses the supplied transaction") + require.Zero(t, execTx.commitCalls, "the caller owns the commit") + require.Zero(t, execTx.rollbackCalls, "the caller owns error recovery") + }) + } +}