Code review feedback to allow legacy hooks/middlewares to work unchanged

Author: brandur

internal/pluginlookup/plugin_lookup.go

@@ -1,6 +1,7 @@
 package pluginlookup
 
 import (
+	"context"
 	"sync"
 
 	"github.com/riverqueue/river/rivertype"
@@ -86,26 +87,28 @@ func NewPluginLookup(plugins []rivertype.Plugin) PluginLookupInterface {
 }
 
 // NormalizePlugins converts hook, middleware, and plugin registrations into a
-// single plugin slice. Hook and middleware registrations must already satisfy
-// rivertype.Plugin to participate.
+// single plugin slice while preserving legacy hook and middleware registrations
+// that don't yet opt into Plugin.
 func NormalizePlugins(hooks []rivertype.Hook, middlewares []rivertype.Middleware, plugins []rivertype.Plugin) []rivertype.Plugin {
 	normalizedPlugins := make([]rivertype.Plugin, 0, len(hooks)+len(middlewares)+len(plugins))
 
 	for _, hook := range hooks {
 		if plugin, ok := hook.(rivertype.Plugin); ok {
 			normalizedPlugins = append(normalizedPlugins, plugin)
+			continue
 		}
+
+		normalizedPlugins = append(normalizedPlugins, newHookPlugin(hook))
 	}
 	for _, middleware := range middlewares {
 		if plugin, ok := middleware.(rivertype.Plugin); ok {
 			normalizedPlugins = append(normalizedPlugins, plugin)
+			continue
 		}
+
+		normalizedPlugins = append(normalizedPlugins, newMiddlewarePlugin(middleware))
 	}
-	for _, plugin := range plugins {
-		if plugin != nil {
-			normalizedPlugins = append(normalizedPlugins, plugin)
-		}
-	}
+	normalizedPlugins = append(normalizedPlugins, plugins...)
 
 	return normalizedPlugins
 }
@@ -141,6 +144,89 @@ func (c *emptyPluginLookup) ByKind(kind PluginKind) []rivertype.Plugin {
 	return nil
 }
 
+//
+// hookPlugin
+//
+
+// hookPlugin wraps a legacy hook so it can participate in plugin lookup.
+type hookPlugin struct {
+	hook rivertype.Hook
+}
+
+func newHookPlugin(hook rivertype.Hook) rivertype.Plugin {
+	return &hookPlugin{hook: hook}
+}
+
+func (p *hookPlugin) IsPlugin() bool { return true }
+func (p *hookPlugin) IsHook() bool   { return true }
+
+func (p *hookPlugin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error {
+	hook, ok := p.hook.(rivertype.HookInsertBegin)
+	if !ok {
+		return nil
+	}
+	return hook.InsertBegin(ctx, params)
+}
+
+func (p *hookPlugin) Start(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error {
+	hook, ok := p.hook.(rivertype.HookPeriodicJobsStart)
+	if !ok {
+		return nil
+	}
+	return hook.Start(ctx, params)
+}
+
+func (p *hookPlugin) WorkBegin(ctx context.Context, job *rivertype.JobRow) error {
+	hook, ok := p.hook.(rivertype.HookWorkBegin)
+	if !ok {
+		return nil
+	}
+	return hook.WorkBegin(ctx, job)
+}
+
+func (p *hookPlugin) WorkEnd(ctx context.Context, job *rivertype.JobRow, err error) error {
+	hook, ok := p.hook.(rivertype.HookWorkEnd)
+	if !ok {
+		return err
+	}
+	return hook.WorkEnd(ctx, job, err)
+}
+
+//
+// middlewarePlugin
+//
+
+// middlewarePlugin wraps a legacy middleware so it can participate in plugin
+// lookup.
+type middlewarePlugin struct {
+	middleware rivertype.Middleware
+}
+
+func newMiddlewarePlugin(middleware rivertype.Middleware) rivertype.Plugin {
+	return &middlewarePlugin{middleware: middleware}
+}
+
+func (p *middlewarePlugin) IsPlugin() bool { return true }
+func (p *middlewarePlugin) IsMiddleware() bool {
+	return true
+}
+
+func (p *middlewarePlugin) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) {
+	middleware, ok := p.middleware.(rivertype.JobInsertMiddleware)
+	if !ok {
+		return doInner(ctx)
+	}
+	return middleware.InsertMany(ctx, manyParams, doInner)
+}
+
+func (p *middlewarePlugin) Work(ctx context.Context, job *rivertype.JobRow, doInner func(context.Context) error) error {
+	middleware, ok := p.middleware.(rivertype.WorkerMiddleware)
+	if !ok {
+		return doInner(ctx)
+	}
+	return middleware.Work(ctx, job, doInner)
+}
+
 //
 // JobPluginLookup
 //

internal/pluginlookup/plugin_lookup_test.go

@@ -204,6 +204,40 @@ func TestPluginLookup(t *testing.T) {
 
 		wg.Wait()
 	})
+
+	t.Run("LooksUpLegacyHooksAndMiddleware", func(t *testing.T) {
+		t.Parallel()
+
+		legacyHook := &testLegacyHookInsertBegin{}
+		legacyMiddleware := &testLegacyMiddlewareJobInsert{}
+
+		lookup, isPluginLookup := NewPluginLookup(NormalizePlugins(
+			[]rivertype.Hook{legacyHook},
+			[]rivertype.Middleware{legacyMiddleware},
+			nil,
+		)).(*pluginLookup)
+		require.True(t, isPluginLookup)
+
+		hookPlugins := lookup.ByKind(HookKindInsertBegin)
+		require.Len(t, hookPlugins, 1)
+
+		hook, isHookInsertBegin := hookPlugins[0].(rivertype.HookInsertBegin)
+		require.True(t, isHookInsertBegin)
+		require.NoError(t, hook.InsertBegin(context.Background(), &rivertype.JobInsertParams{}))
+		require.True(t, legacyHook.insertBeginCalled)
+
+		middlewarePlugins := lookup.ByKind(MiddlewareKindJobInsert)
+		require.Len(t, middlewarePlugins, 1)
+
+		middleware, isJobInsertMiddleware := middlewarePlugins[0].(rivertype.JobInsertMiddleware)
+		require.True(t, isJobInsertMiddleware)
+		_, err := middleware.InsertMany(context.Background(), []*rivertype.JobInsertParams{{}}, func(context.Context) ([]*rivertype.JobInsertResult, error) {
+			legacyMiddleware.doInnerCalled = true
+			return nil, nil
+		})
+		require.NoError(t, err)
+		require.True(t, legacyMiddleware.doInnerCalled)
+	})
 }
 
 //
@@ -349,3 +383,36 @@ func (t *testMiddlewareWorker) IsPlugin() bool { return true }
 func (t *testMiddlewareWorker) Work(ctx context.Context, job *rivertype.JobRow, doInner func(context.Context) error) error {
 	return doInner(ctx)
 }
+
+//
+// testLegacyHookInsertBegin
+//
+
+var _ rivertype.HookInsertBegin = &testLegacyHookInsertBegin{}
+
+type testLegacyHookInsertBegin struct {
+	insertBeginCalled bool
+}
+
+func (t *testLegacyHookInsertBegin) IsHook() bool { return true }
+
+func (t *testLegacyHookInsertBegin) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error {
+	t.insertBeginCalled = true
+	return nil
+}
+
+//
+// testLegacyMiddlewareJobInsert
+//
+
+var _ rivertype.JobInsertMiddleware = &testLegacyMiddlewareJobInsert{}
+
+type testLegacyMiddlewareJobInsert struct {
+	doInnerCalled bool
+}
+
+func (t *testLegacyMiddlewareJobInsert) IsMiddleware() bool { return true }
+
+func (t *testLegacyMiddlewareJobInsert) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) {
+	return doInner(ctx)
+}