Let job args return plugins like they can for hooks/middleware

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Added `JobArgsWithPlugins` for installing plugins that apply only to a specific job type. [PR #1337](https://github.com/riverqueue/river/pull/1337).
+
 ## [0.41.1] - 2026-07-29
 
 ### Fixed

client.go

@@ -241,7 +241,8 @@ type Config struct {
 	// and the work hook between them will not run. When a job is worked, the
 	// work hook runs and the insertion hooks on either side of it are skipped.
 	//
-	// Jobs may have their own specific hooks by implementing JobArgsWithHooks.
+	// Jobs may have their own specific hooks by implementing JobArgsWithHooks or
+	// JobArgsWithPlugins.
 	//
 	// Entries in Hooks are installed only as hooks, even if they also implement
 	// rivertype.Middleware. Use Plugins for an extension that should act as
@@ -297,6 +298,8 @@ type Config struct {
 	//
 	// Use Hooks or Middleware when an extension should be installed only as the
 	// corresponding kind. Use Plugins when it should be eligible as both.
+	// Jobs may have their own specific plugins by implementing
+	// JobArgsWithPlugins.
 	Plugins []rivertype.Plugin
 
 	// PeriodicJobs are a set of periodic jobs to run at the specified intervals
@@ -708,7 +711,7 @@ type Client[TTx any] struct {
 	driver                riverdriver.Driver[TTx]
 	elector               *leadership.Elector
 	pluginLookupByJob     *pluginlookup.JobPluginLookup
-	pluginLookupGlobal    pluginlookup.PluginLookupInterface
+	pluginLookupGlobal    *pluginlookup.PluginLookup
 	insertNotifyLimiter   *notifylimiter.Limiter
 	notifier              *notifier.Notifier // may be nil in poll-only mode
 	periodicJobs          *PeriodicJobBundle
@@ -832,9 +835,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 		middleware = pluginconfig.CombinedMiddleware(config.Middleware, config.JobInsertMiddleware, config.WorkerMiddleware)
 		plugins    = append(riverplugin.DefaultPlugins(), config.Plugins...)
 	)
-	pluginlookup.InitBaseServices(archetype, config.Hooks)
-	pluginlookup.InitBaseServices(archetype, middleware)
-	pluginlookup.InitBaseServices(archetype, plugins)
+	pluginLookupByJob := pluginlookup.NewJobPluginLookup(archetype)
+	pluginLookupGlobal := pluginlookup.NewPluginLookupFromConfig(archetype, config.Hooks, middleware, plugins)
 
 	client := &Client[TTx]{
 		clientNotifyBundle: &ClientNotifyBundle[TTx]{
@@ -843,8 +845,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 		},
 		config:               config,
 		driver:               driver,
-		pluginLookupByJob:    pluginlookup.NewJobPluginLookup(),
-		pluginLookupGlobal:   pluginlookup.NewPluginLookupFromConfig(config.Hooks, middleware, plugins),
+		pluginLookupByJob:    pluginLookupByJob,
+		pluginLookupGlobal:   pluginLookupGlobal,
 		producersByQueueName: make(map[string]*producer),
 		testSignals:          clientTestSignals{},
 		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
 	if config.Workers != nil {
 		workerMetadata = make([]*rivertype.WorkerMetadata, 0, len(config.Workers.workersMap))
 		for kind, workerInfo := range config.Workers.workersMap {
-			var hooks []rivertype.Hook
-			if jobArgsWithHooks, ok := workerInfo.jobArgs.(JobArgsWithHooks); ok {
-				hooks = jobArgsWithHooks.Hooks()
-			}
-
 			workerMetadata = append(workerMetadata, &rivertype.WorkerMetadata{
-				JobArgHooks: hooks,
+				JobArgHooks: pluginLookupByJob.ByJobArgs(workerInfo.jobArgs).Hooks(),
 				Kind:        kind,
 			})
 		}
@@ -2004,7 +2001,20 @@ func (c *Client[TTx]) insertManyShared(
 		return insertResults, nil
 	}
 
-	jobInsertMiddleware := c.pluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareJobInsert)
+	jobInsertMiddleware := append([]any(nil), c.pluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareJobInsert)...)
+	jobKindsSeen := make(map[string]struct{}, len(insertParams))
+	for _, params := range insertParams {
+		kind := params.Args.Kind()
+		if _, ok := jobKindsSeen[kind]; ok {
+			continue
+		}
+		jobKindsSeen[kind] = struct{}{}
+
+		jobInsertMiddleware = append(
+			jobInsertMiddleware,
+			c.pluginLookupByJob.ByJobArgs(params.Args).ByKind(pluginlookup.PluginKindMiddlewareJobInsert)...,
+		)
+	}
 	if len(jobInsertMiddleware) > 0 {
 		// Wrap middlewares in reverse order so the one defined first is wrapped
 		// as the outermost function and is first to receive the operation.

internal/jobexecutor/job_executor.go

@@ -112,7 +112,7 @@ type JobExecutor struct {
 	DefaultClientRetryPolicy ClientRetryPolicy
 	ErrorHandler             ErrorHandler
 	PluginLookupByJob        *pluginlookup.JobPluginLookup
-	PluginLookupGlobal       pluginlookup.PluginLookupInterface
+	PluginLookupGlobal       *pluginlookup.PluginLookup
 	JobRow                   *rivertype.JobRow
 	ProducerCallbacks        struct {
 		JobDone func(jobRow *rivertype.JobRow)
@@ -217,12 +217,13 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) {
 		)
 		return &jobExecutorResult{Err: &rivertype.UnknownJobKindError{Kind: e.JobRow.Kind}, MetadataUpdates: metadataUpdates}
 	}
+	pluginLookupByJob := e.WorkUnit.PluginLookup(e.PluginLookupByJob)
 
 	doInner := execution.Func(func(ctx context.Context) error {
 		{
 			for _, hook := range append(
 				e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindHookWorkBegin),
-				e.WorkUnit.PluginLookup(e.PluginLookupByJob).ByKind(pluginlookup.PluginKindHookWorkBegin)...,
+				pluginLookupByJob.ByKind(pluginlookup.PluginKindHookWorkBegin)...,
 			) {
 				if err := hook.(rivertype.HookWorkBegin).WorkBegin(ctx, e.JobRow); err != nil { //nolint:forcetypeassert
 					return err
@@ -251,7 +252,7 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) {
 		{
 			for _, hook := range append(
 				e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindHookWorkEnd),
-				e.WorkUnit.PluginLookup(e.PluginLookupByJob).ByKind(pluginlookup.PluginKindHookWorkEnd)...,
+				pluginLookupByJob.ByKind(pluginlookup.PluginKindHookWorkEnd)...,
 			) {
 				err = hook.(rivertype.HookWorkEnd).WorkEnd(ctx, e.JobRow, err) //nolint:forcetypeassert
 			}
@@ -260,13 +261,19 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) {
 		return err
 	})
 
-	globalMiddleware := make([]rivertype.Middleware, 0, len(e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker)))
+	pluginMiddleware := make([]rivertype.Middleware, 0,
+		len(e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker))+
+			len(pluginLookupByJob.ByKind(pluginlookup.PluginKindMiddlewareWorker)),
+	)
 	for _, plugin := range e.PluginLookupGlobal.ByKind(pluginlookup.PluginKindMiddlewareWorker) {
-		globalMiddleware = append(globalMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert
+		pluginMiddleware = append(pluginMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert
+	}
+	for _, plugin := range pluginLookupByJob.ByKind(pluginlookup.PluginKindMiddlewareWorker) {
+		pluginMiddleware = append(pluginMiddleware, plugin.(rivertype.Middleware)) //nolint:forcetypeassert
 	}
 
 	executeFunc := execution.MiddlewareChain(
-		globalMiddleware,
+		pluginMiddleware,
 		e.WorkUnit.Middleware(),
 		doInner,
 		e.JobRow,

internal/jobexecutor/job_executor_test.go

@@ -38,7 +38,7 @@ type customizableWorkUnit struct {
 	work         func() error
 }
 
-func (w *customizableWorkUnit) PluginLookup(lookup *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface {
+func (w *customizableWorkUnit) PluginLookup(lookup *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup {
 	return pluginlookup.NewPluginLookup(nil)
 }
 
@@ -189,7 +189,7 @@ func TestJobExecutor_Execute(t *testing.T) {
 			Completer:                bundle.completer,
 			DefaultClientRetryPolicy: &retrypolicytest.RetryPolicyNoJitter{},
 			ErrorHandler:             bundle.errorHandler,
-			PluginLookupByJob:        pluginlookup.NewJobPluginLookup(),
+			PluginLookupByJob:        pluginlookup.NewJobPluginLookup(nil),
 			PluginLookupGlobal:       pluginlookup.NewPluginLookup(nil),
 			JobRow:                   bundle.jobRow,
 			ProducerCallbacks: struct {

internal/maintenance/job_rescuer_test.go

@@ -50,7 +50,7 @@ type callbackWorkUnit struct {
 	unmarshalErr error
 }
 
-func (w *callbackWorkUnit) PluginLookup(cache *pluginlookup.JobPluginLookup) pluginlookup.PluginLookupInterface {
+func (w *callbackWorkUnit) PluginLookup(cache *pluginlookup.JobPluginLookup) *pluginlookup.PluginLookup {
 	return nil
 }
 func (w *callbackWorkUnit) Middleware() []rivertype.WorkerMiddleware { return nil }

internal/maintenance/periodic_job_enqueuer.go

@@ -89,7 +89,7 @@ type InsertFunc func(ctx context.Context, tx riverdriver.ExecutorTx, insertParam
 type PeriodicJobEnqueuerConfig struct {
 	AdvisoryLockPrefix int32
 
-	PluginLookupGlobal pluginlookup.PluginLookupInterface
+	PluginLookupGlobal *pluginlookup.PluginLookup
 
 	// Insert is the function to call to insert jobs into the database.
 	Insert InsertFunc