logging: make WithContext() optional

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

v1/logging/buffered_logger.go

@@ -1,7 +1,6 @@
 package logging
 
 import (
-	"context"
 	"fmt"
 	"maps"
 	"sync"
@@ -102,11 +101,6 @@ func (b *BufferedLogger) WithFields(fields map[string]any) Logger {
 	}
 }
 
-// WithContext returns a new logger with context.
-func (b *BufferedLogger) WithContext(context.Context) Logger {
-	return b
-}
-
 // GetLevel returns the current log level.
 func (b *BufferedLogger) GetLevel() Level {
 	b.mu.Lock()
@@ -179,10 +173,6 @@ func (b *bufferedLoggerWithFields) WithFields(fields map[string]any) Logger {
 	}
 }
 
-func (b *bufferedLoggerWithFields) WithContext(context.Context) Logger {
-	return b
-}
-
 func (b *bufferedLoggerWithFields) GetLevel() Level {
 	return b.parent.GetLevel()
 }

v1/logging/buffered_logger_test.go

@@ -1,7 +1,6 @@
 package logging
 
 import (
-	"context"
 	"fmt"
 	"testing"
 )
@@ -176,10 +175,6 @@ func (c *captureLogger) WithFields(fields map[string]any) Logger {
 	return c
 }
 
-func (c *captureLogger) WithContext(context.Context) Logger {
-	return c
-}
-
 func (c *captureLogger) GetLevel() Level {
 	return c.level
 }

v1/logging/logging.go

@@ -34,12 +34,28 @@ type Logger interface {
 	Warn(fmt string, a ...any)
 
 	WithFields(map[string]any) Logger
-	WithContext(context.Context) Logger
 
 	GetLevel() Level
 	SetLevel(Level)
 }
 
+// LoggerWithContext is an optional interface that Logger implementations
+// can implement to support extracting trace information from a context.
+// Use WithContext to call this method on a Logger if it is supported.
+type LoggerWithContext interface {
+	WithContext(context.Context) Logger
+}
+
+// WithContext returns a logger with context information if the logger
+// supports it (i.e., implements LoggerWithContext). Otherwise, the
+// logger is returned unchanged.
+func WithContext(logger Logger, ctx context.Context) Logger {
+	if lc, ok := logger.(LoggerWithContext); ok {
+		return lc.WithContext(ctx)
+	}
+	return logger
+}
+
 // StandardLogger is the default OPA logger implementation.
 type StandardLogger struct {
 	logger *logrus.Logger
@@ -81,13 +97,6 @@ func (l *StandardLogger) WithFields(fields map[string]any) Logger {
 	return &cp
 }
 
-// WithContext returns a logger with trace information extracted from the context.
-// Currently returns the logger as-is. Logger plugins can override this to extract
-// trace/span IDs or other context-specific information.
-func (l *StandardLogger) WithContext(context.Context) Logger {
-	return l
-}
-
 // getFields returns additional fields of this logger
 func (l *StandardLogger) getFields() map[string]any {
 	return l.fields
@@ -190,11 +199,6 @@ func (l *NoOpLogger) WithFields(fields map[string]any) Logger {
 	return &cp
 }
 
-// WithContext returns the logger unchanged (no-op).
-func (l *NoOpLogger) WithContext(context.Context) Logger {
-	return l
-}
-
 // Debug noop
 func (*NoOpLogger) Debug(string, ...any) {}
 
@@ -326,7 +330,7 @@ func (h *SlogHandler) Handle(ctx context.Context, record slog.Record) error {
 		logger = logger.WithFields(attrs)
 	}
 	if ctx != nil {
-		logger = logger.WithContext(ctx)
+		logger = WithContext(logger, ctx)
 	}
 
 	msg := record.Message
@@ -367,6 +371,8 @@ type loggerFromSlogHandler struct {
 	ctx     context.Context
 }
 
+var _ LoggerWithContext = (*loggerFromSlogHandler)(nil)
+
 // NewLoggerFromSlogHandler creates a Logger from an slog.Handler
 func NewLoggerFromSlogHandler(handler slog.Handler, level Level) Logger {
 	return &loggerFromSlogHandler{

v1/logging/test/test.go

@@ -1,7 +1,6 @@
 package test
 
 import (
-	"context"
 	"fmt"
 	"maps"
 	"sync"
@@ -53,11 +52,6 @@ func (l *Logger) WithFields(fields map[string]any) logging.Logger {
 	return &cp
 }
 
-// WithContext returns the logger unchanged (test logger doesn't use context).
-func (l *Logger) WithContext(context.Context) logging.Logger {
-	return l
-}
-
 // Debug buffers a log message.
 func (l *Logger) Debug(f string, a ...any) {
 	l.append(logging.Debug, f, a...)

v1/plugins/logs/logger_plugin_integration_test.go

@@ -334,10 +334,6 @@ func (l *mockLogger) WithFields(fields map[string]any) logging.Logger {
 	}
 }
 
-func (l *mockLogger) WithContext(context.Context) logging.Logger {
-	return l
-}
-
 func (l *mockLogger) GetLevel() logging.Level {
 	return l.level
 }

v1/runtime/logging.go

@@ -32,7 +32,7 @@ func (h loggingPrintHook) Print(pctx print.Context, msg string) error {
 		fields = make(map[string]any, 1)
 	}
 	fields["line"] = pctx.Location.String()
-	h.logger.WithFields(fields).WithContext(pctx.Context).Info(msg)
+	logging.WithContext(h.logger.WithFields(fields), pctx.Context).Info(msg)
 	return nil
 }
 
@@ -112,7 +112,7 @@ func (h *LoggingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			fields["req_params"] = r.URL.Query()
 		}
 
-		log := h.logger.WithContext(r.Context())
+		log := logging.WithContext(h.logger, r.Context())
 		if err == nil {
 			log.WithFields(fields).Info("Received request.")
 		} else {
@@ -158,10 +158,10 @@ func (h *LoggingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 						defer gzReader.Close()
 						fields["resp_body"] = string(plainOutput)
 					} else {
-						h.logger.WithContext(r.Context()).Error("Failed to decompressed the payload: %v", readErr.Error())
+						logging.WithContext(h.logger, r.Context()).Error("Failed to decompressed the payload: %v", readErr.Error())
 					}
 				} else {
-					h.logger.WithContext(r.Context()).Error("Failed to read the compressed payload: %v", gzErr.Error())
+					logging.WithContext(h.logger, r.Context()).Error("Failed to read the compressed payload: %v", gzErr.Error())
 				}
 
 			default:
@@ -169,7 +169,7 @@ func (h *LoggingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			}
 		}
 
-		h.logger.WithFields(fields).WithContext(r.Context()).Info("Sent response.")
+		logging.WithContext(h.logger.WithFields(fields), r.Context()).Info("Sent response.")
 	}
 }