feat: context attributes

Author: nrwiersma

bench_test.go

@@ -1,6 +1,7 @@
 package logger_test
 
 import (
+	"context"
 	"log/slog"
 	"os"
 	"testing"
@@ -102,6 +103,32 @@ func BenchmarkLogger_JsonCtx(b *testing.B) {
 	})
 }
 
+func BenchmarkLogger_WithContext(b *testing.B) {
+	log := logger.New(discard{}, logger.LogfmtFormat(), logger.Debug)
+	goCtx := context.Background()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			_ = logger.WithContext(goCtx, log, ctx.Str("req_id", "abc123"), ctx.Int("attempt", 1))
+		}
+	})
+}
+
+func BenchmarkLogger_FromContext(b *testing.B) {
+	log := logger.New(discard{}, logger.LogfmtFormat(), logger.Debug).With(ctx.Str("_n", "bench"))
+	goCtx := logger.WithContext(context.Background(), log, ctx.Str("req_id", "abc123"), ctx.Int("attempt", 1))
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			log.FromContext(goCtx).Info("some message")
+		}
+	})
+}
+
 func BenchmarkHandler_Logfmt(b *testing.B) {
 	h := logger.NewHandler(discard{}, logger.LogfmtFormat(), slog.LevelDebug)
 	log := slog.New(h)

context.go

@@ -0,0 +1,68 @@
+package logger
+
+import "context"
+
+type contextKey struct{}
+
+type ctxFields struct {
+	b []byte
+}
+
+// WithContext returns a new context carrying the given fields.
+// Any fields previously attached via WithContext are preserved;
+// the new fields are appended after them. Any fields already
+// added to the logger are not attached to the context.
+//
+// If the logger discards output or no fields are given, ctx is returned
+// unchanged.
+func WithContext(ctx context.Context, log *Logger, fields ...Field) context.Context {
+	if log.isDiscard || len(fields) == 0 {
+		return ctx
+	}
+
+	e := newEvent(log.fmtr)
+	defer putEvent(e)
+
+	if existing, _ := ctx.Value(contextKey{}).(*ctxFields); existing != nil {
+		e.buf.Write(existing.b)
+	}
+
+	for _, field := range fields {
+		field(e)
+	}
+
+	b := make([]byte, e.buf.Len())
+	copy(b, e.buf.Bytes())
+
+	return context.WithValue(ctx, contextKey{}, &ctxFields{b: b})
+}
+
+// FromContext returns a new Logger extended with any fields attached to ctx
+// via WithContext. The context fields appear after the logger's own pre-rendered
+// fields (set via With) and before per-call fields.
+//
+// If the logger discards output or ctx carries no fields, the receiver is
+// returned unchanged.
+func (l *Logger) FromContext(ctx context.Context) *Logger {
+	if l.isDiscard {
+		return l
+	}
+
+	fields, ok := ctx.Value(contextKey{}).(*ctxFields)
+	if !ok || len(fields.b) == 0 {
+		return l
+	}
+
+	b := make([]byte, len(l.ctx)+len(fields.b))
+	copy(b, l.ctx)
+	copy(b[len(l.ctx):], fields.b)
+
+	return &Logger{
+		w:         l.w,
+		isDiscard: l.isDiscard,
+		fmtr:      l.fmtr,
+		timeFn:    l.timeFn,
+		lvl:       l.lvl,
+		ctx:       b,
+	}
+}

context_test.go

@@ -0,0 +1,95 @@
+package logger_test
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"testing"
+
+	"github.com/hamba/logger/v2"
+	"github.com/hamba/logger/v2/ctx"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLogger_WithContextAttachesFields_Logfmt(t *testing.T) {
+	t.Parallel()
+
+	var buf bytes.Buffer
+	log := logger.New(&buf, logger.LogfmtFormat(), logger.Info).With(ctx.Str("svc", "api"))
+
+	goCtx := logger.WithContext(context.Background(), log, ctx.Str("req_id", "abc123"))
+	log.FromContext(goCtx).Info("handled")
+
+	assert.Equal(t, "lvl=info msg=handled svc=api req_id=abc123\n", buf.String())
+}
+
+func TestLogger_WithContextAttachesFields_JSON(t *testing.T) {
+	t.Parallel()
+
+	var buf bytes.Buffer
+	log := logger.New(&buf, logger.JSONFormat(), logger.Info).With(ctx.Str("svc", "api"))
+
+	goCtx := logger.WithContext(context.Background(), log, ctx.Str("req_id", "abc123"))
+	log.FromContext(goCtx).Info("handled")
+
+	assert.Equal(t, `{"lvl":"info","msg":"handled","svc":"api","req_id":"abc123"}`+"\n", buf.String())
+}
+
+func TestLogger_WithContextCanLayersFields(t *testing.T) {
+	t.Parallel()
+
+	var buf bytes.Buffer
+	log := logger.New(&buf, logger.LogfmtFormat(), logger.Info)
+
+	goCtx := logger.WithContext(context.Background(), log, ctx.Str("req_id", "abc123"))
+	goCtx = logger.WithContext(goCtx, log, ctx.Str("user", "u456"))
+	log.FromContext(goCtx).Info("handled")
+
+	assert.Equal(t, "lvl=info msg=handled req_id=abc123 user=u456\n", buf.String())
+}
+
+func TestLogger_WithContextWithNoFieldsReturnsCtxUnchanged(t *testing.T) {
+	t.Parallel()
+
+	log := logger.New(io.Discard, logger.LogfmtFormat(), logger.Info)
+	goCtx := context.Background()
+
+	got := logger.WithContext(goCtx, log)
+
+	assert.Equal(t, goCtx, got)
+}
+
+func TestLogger_WithContextWithDiscardReturnsCtxUnchanged(t *testing.T) {
+	t.Parallel()
+
+	log := logger.New(io.Discard, logger.LogfmtFormat(), logger.Info)
+	goCtx := context.Background()
+
+	got := logger.WithContext(goCtx, log, ctx.Str("req_id", "abc123"))
+
+	assert.Equal(t, goCtx, got)
+}
+
+func TestLogger_FromContextWithNoContextFieldsReturnsSameLogger(t *testing.T) {
+	t.Parallel()
+
+	var buf bytes.Buffer
+	log := logger.New(&buf, logger.LogfmtFormat(), logger.Info)
+
+	got := log.FromContext(context.Background())
+
+	assert.Same(t, log, got)
+}
+
+func TestLogger_FromContextWithDiscardReturnsSameLogger(t *testing.T) {
+	t.Parallel()
+
+	log := logger.New(io.Discard, logger.LogfmtFormat(), logger.Info)
+	goCtx := logger.WithContext(context.Background(), log, ctx.Str("req_id", "abc123"))
+
+	// Re-attach to a non-discard logger to prove the discard check fires first.
+	discardLog := logger.New(io.Discard, logger.LogfmtFormat(), logger.Info)
+	got := discardLog.FromContext(goCtx)
+
+	assert.Same(t, discardLog, got)
+}

example_test.go

@@ -1,6 +1,7 @@
 package logger_test
 
 import (
+	"context"
 	"log/slog"
 	"os"
 
@@ -35,3 +36,12 @@ func ExampleNewHandler_logfmt() {
 
 	log.Info("connected", slog.String("driver", "pgx"))
 }
+
+func ExampleWithContext() {
+	log := logger.New(os.Stdout, logger.LogfmtFormat(), logger.Info).With(ctx.Str("svc", "api"))
+
+	reqCtx := logger.WithContext(context.Background(), log, ctx.Str("req_id", "abc-123"), ctx.Str("method", "GET"))
+
+	reqLog := log.FromContext(reqCtx)
+	reqLog.Info("request handled", ctx.Int("status", 200))
+}