Arbitrary context call for riverlog middleware

This one's inspired by #914. Our current `riverlog` implementation
should work pretty well for most users, but it does force them to use
slog. I'm inclined to provide additional options here because in all
honesty, slog's API is terrible, and I wouldn't want to use it myself.

Here, we add the new initializer `NewMiddlewareArbitrary`. Users will
have to bring their own context key like:

    type ArbitraryContextKey struct{}

New middleware takes a `newContext` function that should initialize
whatever it wants with an input writer and put it in context:

    riverlog.NewMiddlewareArbitrary(func(ctx context.Context, w io.Writer) context.Context {
        logger := log.New(w, "", 0)
        return context.WithValue(ctx, ArbitraryContextKey{}, logger)
    }, nil),

Work functions then re-extract whatever they added to context (in
reality, they'd probably define their own shorthand helpers over what's
shown below):

    func (w *ArbitraryLoggingWorker) Work(ctx context.Context, job *river.Job[ArbitraryLogginArgs]) error {
        logger := ctx.Value(ArbitraryContextKey{}).(*log.Logger)
        logger.Printf("Raw log from worker")
        return nil
    }

An alternative to this might be to expose more River innards that'd let
callers set their own metadata, but I'm kind of thinking that a project
like that should involve very thorough consideration, and may not even
be desirable. I think that even if we go that direction, the new
`NewMiddlewareArbitrary` would continue to be useful as a higher level
abstraction.

Fixes #914.

Author: brandur

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Preliminary River driver for SQLite (`riverdriver/riversqlite`). This driver seems to produce good results as judged by the test suite, but so far has minimal real world vetting. Try it and let us know how it works out. [PR #870](https://github.com/riverqueue/river/pull/870).
 - CLI `river migrate-get` now takes a `--schema` option to inject a custom schema into dumped migrations and schema comments are hidden if `--schema` option isn't provided. [PR #903](https://github.com/riverqueue/river/pull/903).
+- Added `riverlog.NewMiddlewareArbitrary` that makes the use of `riverlog` job-persisted logging possible with non-slog loggers. [PR #919](https://github.com/riverqueue/river/pull/919).
 
 ### Changed
 

riverlog/example_new_middleware_arbitrary_test.go

@@ -0,0 +1,108 @@
+package riverlog_test
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log"
+	"log/slog"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+
+	"github.com/riverqueue/river"
+	"github.com/riverqueue/river/riverdbtest"
+	"github.com/riverqueue/river/riverdriver/riverpgxv5"
+	"github.com/riverqueue/river/riverlog"
+	"github.com/riverqueue/river/rivershared/riversharedtest"
+	"github.com/riverqueue/river/rivershared/util/slogutil"
+	"github.com/riverqueue/river/rivershared/util/testutil"
+	"github.com/riverqueue/river/rivertype"
+)
+
+// Callers should define their own context key to extract their a logger back
+// out of work context.
+type ArbitraryContextKey struct{}
+
+type ArbitraryLogginArgs struct{}
+
+func (ArbitraryLogginArgs) Kind() string { return "logging" }
+
+type ArbitraryLoggingWorker struct {
+	river.WorkerDefaults[ArbitraryLogginArgs]
+}
+
+func (w *ArbitraryLoggingWorker) Work(ctx context.Context, job *river.Job[ArbitraryLogginArgs]) error {
+	logger := ctx.Value(ArbitraryContextKey{}).(*log.Logger) //nolint:forcetypeassert
+	logger.Printf("Raw log from worker")
+	return nil
+}
+
+// ExampleNewMiddlewareArbitrary demonstrates the use of riverlog middleware
+// with an arbitrary new context function that can be used to inject any sort of
+// logger into context.
+func ExampleNewMiddlewareArbitrary() {
+	ctx := context.Background()
+
+	dbPool, err := pgxpool.New(ctx, riversharedtest.TestDatabaseURL())
+	if err != nil {
+		panic(err)
+	}
+	defer dbPool.Close()
+
+	workers := river.NewWorkers()
+	river.AddWorker(workers, &ArbitraryLoggingWorker{})
+
+	riverClient, err := river.NewClient(riverpgxv5.New(dbPool), &river.Config{
+		Logger: slog.New(&slogutil.SlogMessageOnlyHandler{Level: slog.LevelWarn}),
+		Queues: map[string]river.QueueConfig{
+			river.QueueDefault: {MaxWorkers: 100},
+		},
+		Middleware: []rivertype.Middleware{
+			riverlog.NewMiddlewareArbitrary(func(ctx context.Context, w io.Writer) context.Context {
+				// For demonstration purposes we show the use of a built-in
+				// non-slog logger, but this could be anything like Logrus or
+				// Zap. Even the raw writer could be stored if so desired.
+				logger := log.New(w, "", 0)
+				return context.WithValue(ctx, ArbitraryContextKey{}, logger)
+			}, nil),
+		},
+		Schema:   riverdbtest.TestSchema(ctx, testutil.PanicTB(), riverpgxv5.New(dbPool), nil), // only necessary for the example test
+		TestOnly: true,                                                                         // suitable only for use in tests; remove for live environments
+		Workers:  workers,
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	// Out of example scope, but used to wait until a job is worked.
+	subscribeChan, subscribeCancel := riverClient.Subscribe(river.EventKindJobCompleted)
+	defer subscribeCancel()
+
+	if err := riverClient.Start(ctx); err != nil {
+		panic(err)
+	}
+
+	_, err = riverClient.Insert(ctx, ArbitraryLogginArgs{}, nil)
+	if err != nil {
+		panic(err)
+	}
+
+	// Wait for job to complete, extract log data out of metadata, and print it.
+	for _, event := range riversharedtest.WaitOrTimeoutN(testutil.PanicTB(), subscribeChan, 1) {
+		var metadataWithLog metadataWithLog
+		if err := json.Unmarshal(event.Job.Metadata, &metadataWithLog); err != nil {
+			panic(err)
+		}
+		for _, logAttempt := range metadataWithLog.RiverLog {
+			fmt.Print(logAttempt.Log)
+		}
+	}
+
+	if err := riverClient.Stop(ctx); err != nil {
+		panic(err)
+	}
+
+	// Output:
+	// Raw log from worker
+}

riverlog/example_middleware_test.go riverlog/example_new_middleware_test.goriverlog/example_middleware_test.go renamed to riverlog/example_new_middleware_test.go

@@ -33,9 +33,9 @@ func (w *LoggingWorker) Work(ctx context.Context, job *river.Job[LoggingArgs]) e
 	return nil
 }
 
-// Example_middleware demonstrates the use of riverlog middleware to inject a
+// ExampleNewMiddleware demonstrates the use of riverlog middleware to inject a
 // logger into context that'll persist its output onto the job record.
-func Example_middleware() {
+func ExampleNewMiddleware() {
 	ctx := context.Background()
 
 	dbPool, err := pgxpool.New(ctx, riversharedtest.TestDatabaseURL())
@@ -83,12 +83,6 @@ func Example_middleware() {
 		panic(err)
 	}
 
-	type metadataWithLog struct {
-		RiverLog []struct {
-			Log string `json:"log"`
-		} `json:"river:log"`
-	}
-
 	// Wait for job to complete, extract log data out of metadata, and print it.
 	for _, event := range riversharedtest.WaitOrTimeoutN(testutil.PanicTB(), subscribeChan, 1) {
 		var metadataWithLog metadataWithLog
@@ -108,3 +102,9 @@ func Example_middleware() {
 	// Logged from worker
 	// Another line logged from worker
 }
+
+type metadataWithLog struct {
+	RiverLog []struct {
+		Log string `json:"log"`
+	} `json:"river:log"`
+}

riverlog/river_log.go

@@ -42,8 +42,9 @@ func Logger(ctx context.Context) *slog.Logger {
 type Middleware struct {
 	baseservice.BaseService
 	rivertype.Middleware
-	config     *MiddlewareConfig
-	newHandler func(w io.Writer) slog.Handler
+	config           *MiddlewareConfig
+	newContextLogger func(ctx context.Context, w io.Writer) context.Context
+	newSlogHandler   func(w io.Writer) slog.Handler
 }
 
 // MiddlewareConfig is configuration for Middleware.
@@ -64,8 +65,8 @@ type MiddlewareConfig struct {
 	MaxSizeBytes int
 }
 
-// NewMiddleware initializes a new Middleware with the given handler function
-// and configuration.
+// NewMiddleware initializes a new Middleware with the given slog handler
+// initialization function and configuration.
 //
 // newHandler is a function which is invoked on every Work execution to generate
 // a new slog.Handler for a work-specific slog.Logger. It should take an
@@ -77,20 +78,46 @@ type MiddlewareConfig struct {
 //	riverlog.NewMiddleware(func(w io.Writer) slog.Handler {
 //		return slog.NewJSONHandler(w, nil)
 //	}, nil)
-func NewMiddleware(newHandler func(w io.Writer) slog.Handler, config *MiddlewareConfig) *Middleware {
-	if config == nil {
-		config = &MiddlewareConfig{}
+func NewMiddleware(newSlogHandler func(w io.Writer) slog.Handler, config *MiddlewareConfig) *Middleware {
+	return &Middleware{
+		config:         defaultConfig(config),
+		newSlogHandler: newSlogHandler,
 	}
+}
 
-	// Assign defaults.
-	config = &MiddlewareConfig{
-		MaxSizeBytes: cmp.Or(config.MaxSizeBytes, maxSizeBytes),
+// NewMiddlewareArbitrary initializes a new Middleware with the given arbitrary
+// context initialization function and configuration.
+//
+// newContext is a function which is invoked on every Work execution to generate
+// a new context for the worker. It's generally used to initialize a logger with
+// the given writer and put it in context under a user-defined context key for
+// later use.
+//
+// This variant is meant to provide callers with a version of the middleware
+// that's not tied to slog. A non-slog standard library logger, Logrus, or Zap
+// logger could all be placed in context according to preferred convention.
+//
+// For example:
+//
+//	riverlog.NewMiddlewareArbitrary(func(ctx context.Context, w io.Writer) context.Context {
+//		logger := log.New(w, "", 0)
+//		return context.WithValue(ctx, ArbitraryContextKey{}, logger)
+//	}, nil),
+func NewMiddlewareArbitrary(newContext func(ctx context.Context, w io.Writer) context.Context, config *MiddlewareConfig) *Middleware {
+	return &Middleware{
+		config:           defaultConfig(config),
+		newContextLogger: newContext,
 	}
+}
 
-	return &Middleware{
-		config:     config,
-		newHandler: newHandler,
+func defaultConfig(config *MiddlewareConfig) *MiddlewareConfig {
+	if config == nil {
+		config = &MiddlewareConfig{}
 	}
+
+	config.MaxSizeBytes = cmp.Or(config.MaxSizeBytes, maxSizeBytes)
+
+	return config
 }
 
 type logAttempt struct {
@@ -106,9 +133,18 @@ func (m *Middleware) Work(ctx context.Context, job *rivertype.JobRow, doInner fu
 	var (
 		existingLogData metadataWithLog
 		logBuf          bytes.Buffer
-		logger          = slog.New(m.newHandler(&logBuf))
 	)
 
+	switch {
+	case m.newContextLogger != nil:
+		ctx = m.newContextLogger(ctx, &logBuf)
+	case m.newSlogHandler != nil:
+		logger := slog.New(m.newSlogHandler(&logBuf))
+		ctx = context.WithValue(ctx, contextKey{}, logger)
+	default:
+		return errors.New("expected either newContextLogger or newSlogHandler to be set")
+	}
+
 	if err := json.Unmarshal(job.Metadata, &existingLogData); err != nil {
 		return err
 	}
@@ -145,5 +181,5 @@ func (m *Middleware) Work(ctx context.Context, job *rivertype.JobRow, doInner fu
 		metadataUpdates[metadataKey] = json.RawMessage(allLogDataBytes)
 	}()
 
-	return doInner(context.WithValue(ctx, contextKey{}, logger))
+	return doInner(ctx)
 }

riverlog/river_log_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"io"
+	"log"
 	"log/slog"
 	"testing"
 
@@ -54,8 +55,10 @@ func TestMiddleware(t *testing.T) {
 	ctx := context.Background()
 
 	type testBundle struct {
-		driver *riverpgxv5.Driver
-		tx     pgx.Tx
+		clientConfig *river.Config
+		driver       *riverpgxv5.Driver
+		middleware   *Middleware
+		tx           pgx.Tx
 	}
 
 	setup := func(t *testing.T, config *MiddlewareConfig) (*rivertest.Worker[loggingArgs, pgx.Tx], *testBundle) {
@@ -74,8 +77,10 @@ func TestMiddleware(t *testing.T) {
 		)
 
 		return rivertest.NewWorker(t, driver, clientConfig, worker), &testBundle{
-			driver: driver,
-			tx:     tx,
+			clientConfig: clientConfig,
+			driver:       driver,
+			middleware:   middleware,
+			tx:           tx,
 		}
 	}
 
@@ -205,4 +210,39 @@ func TestMiddleware(t *testing.T) {
 			metadataWithLog.RiverLog,
 		)
 	})
+
+	t.Run("RawMiddleware", func(t *testing.T) {
+		t.Parallel()
+
+		_, bundle := setup(t, nil)
+
+		type writerContextKey struct{}
+
+		bundle.middleware.newContextLogger = func(ctx context.Context, w io.Writer) context.Context {
+			logger := log.New(w, "", 0)
+			return context.WithValue(ctx, writerContextKey{}, logger)
+		}
+		bundle.middleware.newSlogHandler = nil
+
+		testWorker := rivertest.NewWorker(t, bundle.driver, bundle.clientConfig, river.WorkFunc(func(ctx context.Context, job *river.Job[loggingArgs]) error {
+			logger := ctx.Value(writerContextKey{}).(*log.Logger) //nolint:forcetypeassert
+			logger.Printf(job.Args.Message)
+			return nil
+		}))
+
+		workRes, err := testWorker.Work(ctx, t, bundle.tx, loggingArgs{Message: "Raw log from worker"}, nil)
+		require.NoError(t, err)
+
+		var metadataWithLog metadataWithLog
+		require.NoError(t, json.Unmarshal(workRes.Job.Metadata, &metadataWithLog))
+
+		require.Equal(t, []logAttempt{
+			{
+				Attempt: 1,
+				Log:     "Raw log from worker\n",
+			},
+		},
+			metadataWithLog.RiverLog,
+		)
+	})
 }