Skip to content

Commit eba5b6a

Browse files
committed
Let job args return plugins like they can for hooks/middleware
Here, implement a job args `jobArgsWithPlugins` interface that lets job args return a set of `Plugins`. This follows up the addition of plugins in #1284, and makes the functionality symmetrical to what's available with hooks/middleware. I would have done this before, but I forgot that you could do this until I was looking at docs today. As in #1284 for client config, `Hooks` and `Middleware` continue to be supported on job args and have not yet been deprecated. type EmailArgs struct { To string `json:"to"` } func (EmailArgs) Kind() string { return "email" } func (EmailArgs) Plugins() []rivertype.Plugin { return []rivertype.Plugin{ &EmailPlugin{}, } } type EmailPlugin struct { river.PluginDefaults } func (p *EmailPlugin) WorkBegin(ctx context.Context, job *rivertype.JobRow) error { return nil } I also had Codex look for opportunities for implementation clean up since all of this was feeling like a few too many LOCs. It was able to clean up ~85 LOCs net, with the main change being to get rid of `PluginLookupInterface` + `emptyPluginLookup`. This removes a layer of indirection, which is good, but also `emptyPluginLookup` had become less useful due to River installing default middleware (`ResumableMiddleware`) on all clients whether any has been specified by a client or not.
1 parent 0780aa9 commit eba5b6a

17 files changed

Lines changed: 348 additions & 180 deletions

CHANGELOG.md

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

88
## [Unreleased]
99

10+
### Added
11+
12+
- Added `JobArgsWithPlugins` for installing plugins that apply only to a specific job type. [PR #1337](https://github.com/riverqueue/river/pull/1337).
13+
1014
## [0.41.1] - 2026-07-29
1115

1216
### Fixed

client.go

Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -241,7 +241,8 @@ type Config struct {
241241
// and the work hook between them will not run. When a job is worked, the
242242
// work hook runs and the insertion hooks on either side of it are skipped.
243243
//
244-
// Jobs may have their own specific hooks by implementing JobArgsWithHooks.
244+
// Jobs may have their own specific hooks by implementing JobArgsWithHooks or
245+
// JobArgsWithPlugins.
245246
//
246247
// Entries in Hooks are installed only as hooks, even if they also implement
247248
// rivertype.Middleware. Use Plugins for an extension that should act as
@@ -297,6 +298,8 @@ type Config struct {
297298
//
298299
// Use Hooks or Middleware when an extension should be installed only as the
299300
// corresponding kind. Use Plugins when it should be eligible as both.
301+
// Jobs may have their own specific plugins by implementing
302+
// JobArgsWithPlugins.
300303
Plugins []rivertype.Plugin
301304

302305
// PeriodicJobs are a set of periodic jobs to run at the specified intervals
@@ -708,7 +711,7 @@ type Client[TTx any] struct {
708711
driver riverdriver.Driver[TTx]
709712
elector *leadership.Elector
710713
pluginLookupByJob *pluginlookup.JobPluginLookup
711-
pluginLookupGlobal pluginlookup.PluginLookupInterface
714+
pluginLookupGlobal *pluginlookup.PluginLookup
712715
insertNotifyLimiter *notifylimiter.Limiter
713716
notifier *notifier.Notifier // may be nil in poll-only mode
714717
periodicJobs *PeriodicJobBundle
@@ -832,9 +835,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
832835
middleware = pluginconfig.CombinedMiddleware(config.Middleware, config.JobInsertMiddleware, config.WorkerMiddleware)
833836
plugins = append(riverplugin.DefaultPlugins(), config.Plugins...)
834837
)
835-
pluginlookup.InitBaseServices(archetype, config.Hooks)
836-
pluginlookup.InitBaseServices(archetype, middleware)
837-
pluginlookup.InitBaseServices(archetype, plugins)
838+
pluginLookupByJob := pluginlookup.NewJobPluginLookup(archetype)
839+
pluginLookupGlobal := pluginlookup.NewPluginLookupFromConfig(archetype, config.Hooks, middleware, plugins)
838840

839841
client := &Client[TTx]{
840842
clientNotifyBundle: &ClientNotifyBundle[TTx]{
@@ -843,8 +845,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
843845
},
844846
config: config,
845847
driver: driver,
846-
pluginLookupByJob: pluginlookup.NewJobPluginLookup(),
847-
pluginLookupGlobal: pluginlookup.NewPluginLookupFromConfig(config.Hooks, middleware, plugins),
848+
pluginLookupByJob: pluginLookupByJob,
849+
pluginLookupGlobal: pluginLookupGlobal,
848850
producersByQueueName: make(map[string]*producer),
849851
testSignals: clientTestSignals{},
850852
workCancel: func(cause error) {}, // replaced on start, but here in case StopAndCancel is called before start up
@@ -872,13 +874,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
872874
if config.Workers != nil {
873875
workerMetadata = make([]*rivertype.WorkerMetadata, 0, len(config.Workers.workersMap))
874876
for kind, workerInfo := range config.Workers.workersMap {
875-
var hooks []rivertype.Hook
876-
if jobArgsWithHooks, ok := workerInfo.jobArgs.(JobArgsWithHooks); ok {
877-
hooks = jobArgsWithHooks.Hooks()
878-
}
879-
880877
workerMetadata = append(workerMetadata, &rivertype.WorkerMetadata{
881-
JobArgHooks: hooks,
878+
JobArgHooks: pluginLookupByJob.ByJobArgs(workerInfo.jobArgs).Hooks(),
882879
Kind: kind,
883880
})
884881
}
@@ -2004,7 +2001,20 @@ func (c *Client[TTx]) insertManyShared(
20042001
return insertResults, nil
20052002
}
20062003

2007-
jobInsertMiddleware := c.pluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareJobInsert)
2004+
jobInsertMiddleware := append([]any(nil), c.pluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareJobInsert)...)
2005+
jobKindsSeen := make(map[string]struct{}, len(insertParams))
2006+
for _, params := range insertParams {
2007+
kind := params.Args.Kind()
2008+
if _, ok := jobKindsSeen[kind]; ok {
2009+
continue
2010+
}
2011+
jobKindsSeen[kind] = struct{}{}
2012+
2013+
jobInsertMiddleware = append(
2014+
jobInsertMiddleware,
2015+
c.pluginLookupByJob.ByJobArgs(params.Args).ByKind(pluginlookup.PluginKindMiddlewareJobInsert)...,
2016+
)
2017+
}
20082018
if len(jobInsertMiddleware) > 0 {
20092019
// Wrap middlewares in reverse order so the one defined first is wrapped
20102020
// as the outermost function and is first to receive the operation.

internal/jobexecutor/job_executor.go

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ type JobExecutor struct {
112112
DefaultClientRetryPolicy ClientRetryPolicy
113113
ErrorHandler ErrorHandler
114114
PluginLookupByJob *pluginlookup.JobPluginLookup
115-
PluginLookupGlobal pluginlookup.PluginLookupInterface
115+
PluginLookupGlobal *pluginlookup.PluginLookup
116116
JobRow *rivertype.JobRow
117117
ProducerCallbacks struct {
118118
JobDone func(jobRow *rivertype.JobRow)
@@ -217,12 +217,13 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) {
217217
)
218218
return &jobExecutorResult{Err: &rivertype.UnknownJobKindError{Kind: e.JobRow.Kind}, MetadataUpdates: metadataUpdates}
219219
}
220+
pluginLookupByJob := e.WorkUnit.PluginLookup(e.PluginLookupByJob)
220221

221222
doInner := execution.Func(func(ctx context.Context) error {
222223
{
223224
for _, hook := range append(
224225
e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindHookWorkBegin),
225-
e.WorkUnit.PluginLookup(e.PluginLookupByJob).ByKind(pluginlookup.PluginKindHookWorkBegin)...,
226+
pluginLookupByJob.ByKind(pluginlookup.PluginKindHookWorkBegin)...,
226227
) {
227228
if err := hook.(rivertype.HookWorkBegin).WorkBegin(ctx, e.JobRow); err != nil { //nolint:forcetypeassert
228229
return err
@@ -251,7 +252,7 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) {
251252
{
252253
for _, hook := range append(
253254
e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindHookWorkEnd),
254-
e.WorkUnit.PluginLookup(e.PluginLookupByJob).ByKind(pluginlookup.PluginKindHookWorkEnd)...,
255+
pluginLookupByJob.ByKind(pluginlookup.PluginKindHookWorkEnd)...,
255256
) {
256257
err = hook.(rivertype.HookWorkEnd).WorkEnd(ctx, e.JobRow, err) //nolint:forcetypeassert
257258
}
@@ -260,13 +261,19 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) {
260261
return err
261262
})
262263

263-
globalMiddleware := make([]rivertype.Middleware, 0, len(e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker)))
264+
pluginMiddleware := make([]rivertype.Middleware, 0,
265+
len(e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker))+
266+
len(pluginLookupByJob.ByKind(pluginlookup.PluginKindMiddlewareWorker)),
267+
)
264268
for _, plugin := range e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker) {
265-
globalMiddleware = append(globalMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert
269+
pluginMiddleware = append(pluginMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert
270+
}
271+
for _, plugin := range pluginLookupByJob.ByKind(pluginlookup.PluginKindMiddlewareWorker) {
272+
pluginMiddleware = append(pluginMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert
266273
}
267274

268275
executeFunc := execution.MiddlewareChain(
269-
globalMiddleware,
276+
pluginMiddleware,
270277
e.WorkUnit.Middleware(),
271278
doInner,
272279
e.JobRow,

internal/jobexecutor/job_executor_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ type customizableWorkUnit struct {
3838
work func() error
3939
}
4040

41-
func (w *customizableWorkUnit) PluginLookup(lookup *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface {
41+
func (w *customizableWorkUnit) PluginLookup(lookup *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup {
4242
return pluginlookup.NewPluginLookup(nil)
4343
}
4444

@@ -189,7 +189,7 @@ func TestJobExecutor_Execute(t *testing.T) {
189189
Completer: bundle.completer,
190190
DefaultClientRetryPolicy: &retrypolicytest.RetryPolicyNoJitter{},
191191
ErrorHandler: bundle.errorHandler,
192-
PluginLookupByJob: pluginlookup.NewJobPluginLookup(),
192+
PluginLookupByJob: pluginlookup.NewJobPluginLookup(nil),
193193
PluginLookupGlobal: pluginlookup.NewPluginLookup(nil),
194194
JobRow: bundle.jobRow,
195195
ProducerCallbacks: struct {

internal/maintenance/job_rescuer_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ type callbackWorkUnit struct {
5050
unmarshalErr error
5151
}
5252

53-
func (w *callbackWorkUnit) PluginLookup(cache *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface {
53+
func (w *callbackWorkUnit) PluginLookup(cache *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup {
5454
return nil
5555
}
5656
func (w *callbackWorkUnit) Middleware() []rivertype.WorkerMiddleware { return nil }

internal/maintenance/periodic_job_enqueuer.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ type InsertFunc func(ctx context.Context, tx riverdriver.ExecutorTx, insertParam
8989
type PeriodicJobEnqueuerConfig struct {
9090
AdvisoryLockPrefix int32
9191

92-
PluginLookupGlobal pluginlookup.PluginLookupInterface
92+
PluginLookupGlobal *pluginlookup.PluginLookup
9393

9494
// Insert is the function to call to insert jobs into the database.
9595
Insert InsertFunc

0 commit comments

Comments
 (0)