From dd6aa8712486774cf96cea1f97035a48226dd664 Mon Sep 17 00:00:00 2001
From: Patrick Ohly <patrick.ohly@intel.com>
Date: Sun, 27 Aug 2023 22:15:11 +0200
Subject: [PATCH] slogr: add context support

The normal logr API for storing and retrieving a logger is used.
This is necessary to ensure interoperability.

The downside is that storing a slog.Handler and retrieving it again is more
costly than storing and retrieving a logr.Logger, because additional
allocations are needed.
---
 slogr/slogr.go      | 21 +++++++++++++++++++++
 slogr/slogr_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/slogr/slogr.go b/slogr/slogr.go
index 8828fd2..437520e 100644
--- a/slogr/slogr.go
+++ b/slogr/slogr.go
@@ -105,3 +105,24 @@ type SlogSink interface {
 	WithAttrs(attrs []slog.Attr) SlogSink
 	WithGroup(name string) SlogSink
 }
+
+// HandlerFromContext retrieves a slog.Handler from the context or, if no
+// logger is stored there, returns the slog default handler.
+//
+// It uses logr.FromContext and thus is interoperable with code that only knows
+// about the logr API.
+func HandlerFromContext(ctx context.Context) slog.Handler {
+	logger, err := logr.FromContext(ctx)
+	if err == nil {
+		return NewSlogHandler(logger)
+	}
+	return slog.Default().Handler()
+}
+
+// ContextWithHandler creates a new context which contains the given handler.
+//
+// It stores the handler after conversion to a logr.Logger with logr.NewContext
+// and thus is interoperable with code that only knows about the logr API.
+func ContextWithHandler(ctx context.Context, handler slog.Handler) context.Context {
+	return logr.NewContext(ctx, NewLogr(handler))
+}
diff --git a/slogr/slogr_test.go b/slogr/slogr_test.go
index 2e5f0ff..0fbd80f 100644
--- a/slogr/slogr_test.go
+++ b/slogr/slogr_test.go
@@ -254,6 +254,31 @@ func TestConversion(t *testing.T) {
 	}
 }
 
+func TestContext(t *testing.T) {
+	defaultLogger := slog.Default()
+	defer slog.SetDefault(defaultLogger)
+
+	// Retrieve default.
+	expectEqual(t, defaultLogger.Handler(), slogr.HandlerFromContext(context.Background()))
+	discard := slogr.NewSlogHandler(logr.Discard())
+
+	// Retrieve modified default.
+	slog.SetDefault(slog.New(discard))
+	expectEqual(t, discard, slogr.HandlerFromContext(context.Background()))
+
+	// Store and retrieve funcr.
+	funcrLogger := funcr.New(func(_, _ string) {}, funcr.Options{})
+	ctx := logr.NewContext(context.Background(), funcrLogger)
+	funcrHandler := slogr.HandlerFromContext(ctx)
+	expectEqual(t, funcrLogger, slogr.NewLogr(funcrHandler))
+
+	// Store and retrieve slog.JSONHandler.
+	jsonHandler := slog.NewJSONHandler(io.Discard, nil)
+	ctx = slogr.ContextWithHandler(context.Background(), jsonHandler)
+	jsonLogger := logr.FromContextOrDiscard(ctx)
+	expectEqual(t, jsonHandler, slogr.NewSlogHandler(jsonLogger))
+}
+
 func expectEqual(t *testing.T, expected, actual any) {
 	if expected != actual {
 		t.Helper()