contrib/log/slog: fix WithAttrs and WithGroup implementation

contrib/log/slog/slog.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"io"
 	"log/slog"
+	"strconv"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
@@ -23,6 +24,8 @@ func init() {
 	tracer.MarkIntegrationImported("log/slog")
 }
 
+var _ slog.Handler = (*handler)(nil)
+
 // NewJSONHandler is a convenience function that returns a *slog.JSONHandler logger enhanced with
 // tracing information.
 func NewJSONHandler(w io.Writer, opts *slog.HandlerOptions) slog.Handler {
@@ -31,21 +34,83 @@ func NewJSONHandler(w io.Writer, opts *slog.HandlerOptions) slog.Handler {
 
 // WrapHandler enhances the given logger handler attaching tracing information to logs.
 func WrapHandler(h slog.Handler) slog.Handler {
-	return &handler{h}
+	return &handler{wrapped: h}
 }
 
 type handler struct {
-	slog.Handler
+	wrapped    slog.Handler
+	groups     []string
+	groupAttrs map[string][]slog.Attr
+}
+
+// Enabled calls the wrapped handler Enabled method.
+func (h *handler) Enabled(ctx context.Context, level slog.Level) bool {
+	return h.wrapped.Enabled(ctx, level)
 }
 
 // Handle handles the given Record, attaching tracing information if found.
 func (h *handler) Handle(ctx context.Context, rec slog.Record) error {
+	reqHandler := h.wrapped
+
+	// We need to ensure the trace id and span id keys are set at the root level:
+	// https://docs.datadoghq.com/tracing/other_telemetry/connect_logs_and_traces/
+	// In case the user has created group loggers, we ignore those and
+	// set them at the root level.
 	span, ok := tracer.SpanFromContext(ctx)
 	if ok {
-		rec.Add(
-			slog.Uint64(ext.LogKeyTraceID, span.Context().TraceID()),
-			slog.Uint64(ext.LogKeySpanID, span.Context().SpanID()),
-		)
+		traceID := strconv.FormatUint(span.Context().TraceID(), 10)
+		spanID := strconv.FormatUint(span.Context().SpanID(), 10)
+
+		attrs := []slog.Attr{
+			slog.String(ext.LogKeyTraceID, traceID),
+			slog.String(ext.LogKeySpanID, spanID),
+		}
+		reqHandler = reqHandler.WithAttrs(attrs)
+	}
+	for _, group := range h.groups {
+		reqHandler = reqHandler.WithGroup(group)
+		if attrs, ok := h.groupAttrs[group]; ok {
+			reqHandler = reqHandler.WithAttrs(attrs)
+		}
+	}
+	return reqHandler.Handle(ctx, rec)
+}
+
+// WithAttrs saves the provided attributes associated to the current Group.
+// If Group was not called for the logger, we just call WithAttrs for the wrapped handler.
+func (h *handler) WithAttrs(attrs []slog.Attr) slog.Handler {
+	if len(h.groups) == 0 {
+		return &handler{
+			wrapped:    h.wrapped.WithAttrs(attrs),
+			groupAttrs: h.groupAttrs,
+			groups:     h.groups,
+		}
+	}
+	curGroup := h.groups[len(h.groups)-1]
+
+	groupAttrs := groupAttrsCopy(h.groupAttrs)
+	groupAttrs[curGroup] = append(groupAttrs[curGroup], attrs...)
+
+	return &handler{
+		wrapped:    h.wrapped,
+		groupAttrs: groupAttrs,
+		groups:     h.groups,
+	}
+}
+
+// WithGroup saves the provided group to be used later in the Handle method.
+func (h *handler) WithGroup(name string) slog.Handler {
+	return &handler{
+		wrapped:    h.wrapped,
+		groupAttrs: h.groupAttrs,
+		groups:     append(h.groups, name),
+	}
+}
+
+func groupAttrsCopy(m map[string][]slog.Attr) map[string][]slog.Attr {
+	cp := make(map[string][]slog.Attr)
+	for k, v := range m {
+		cp[k] = append([]slog.Attr{}, v...)
 	}
-	return h.Handler.Handle(ctx, rec)
+	return cp
 }

contrib/log/slog/slog_test.go

@@ -9,7 +9,9 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"io"
 	"log/slog"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -21,29 +23,37 @@ import (
 	internallog "gopkg.in/DataDog/dd-trace-go.v1/internal/log"
 )
 
-func assertLogEntry(t *testing.T, rawEntry, wantMsg, wantLevel string) {
+func assertLogEntry(t *testing.T, rawEntry, wantMsg, wantLevel string, span tracer.Span, assertExtra func(t *testing.T, entry map[string]interface{})) {
 	t.Helper()
 
-	var data map[string]interface{}
-	err := json.Unmarshal([]byte(rawEntry), &data)
+	t.Log(rawEntry)
+
+	var entry map[string]interface{}
+	err := json.Unmarshal([]byte(rawEntry), &entry)
 	require.NoError(t, err)
-	require.NotEmpty(t, data)
+	require.NotEmpty(t, entry)
+
+	assert.Equal(t, wantMsg, entry["msg"])
+	assert.Equal(t, wantLevel, entry["level"])
+	assert.NotEmpty(t, entry["time"])
 
-	assert.Equal(t, wantMsg, data["msg"])
-	assert.Equal(t, wantLevel, data["level"])
-	assert.NotEmpty(t, data["time"])
-	assert.NotEmpty(t, data[ext.LogKeyTraceID])
-	assert.NotEmpty(t, data[ext.LogKeySpanID])
+	traceID := strconv.FormatUint(span.Context().TraceID(), 10)
+	spanID := strconv.FormatUint(span.Context().SpanID(), 10)
+	assert.Equal(t, traceID, entry[ext.LogKeyTraceID], "trace id not found")
+	assert.Equal(t, spanID, entry[ext.LogKeySpanID], "span id not found")
+
+	if assertExtra != nil {
+		assertExtra(t, entry)
+	}
 }
 
-func testLogger(t *testing.T, createHandler func(b *bytes.Buffer) slog.Handler) {
+func testLogger(t *testing.T, createLogger func(b io.Writer) *slog.Logger, assertExtra func(t *testing.T, entry map[string]interface{})) {
 	tracer.Start(tracer.WithLogger(internallog.DiscardLogger{}))
 	defer tracer.Stop()
 
 	// create the application logger
 	var b bytes.Buffer
-	h := createHandler(&b)
-	logger := slog.New(h)
+	logger := createLogger(&b)
 
 	// start a new span
 	span, ctx := tracer.StartSpanFromContext(context.Background(), "test")
@@ -59,18 +69,82 @@ func testLogger(t *testing.T, createHandler func(b *bytes.Buffer) slog.Handler)
 	)
 	// assert log entries contain trace information
 	require.Len(t, logs, 2)
-	assertLogEntry(t, logs[0], "this is an info log with tracing information", "INFO")
-	assertLogEntry(t, logs[1], "this is an error log with tracing information", "ERROR")
+	assertLogEntry(t, logs[0], "this is an info log with tracing information", "INFO", span, assertExtra)
+	assertLogEntry(t, logs[1], "this is an error log with tracing information", "ERROR", span, assertExtra)
 }
 
 func TestNewJSONHandler(t *testing.T) {
-	testLogger(t, func(b *bytes.Buffer) slog.Handler {
-		return NewJSONHandler(b, nil)
-	})
+	testLogger(
+		t,
+		func(w io.Writer) *slog.Logger {
+			return slog.New(NewJSONHandler(w, nil))
+		},
+		nil,
+	)
 }
 
 func TestWrapHandler(t *testing.T) {
-	testLogger(t, func(b *bytes.Buffer) slog.Handler {
-		return WrapHandler(slog.NewJSONHandler(b, nil))
+	testLogger(
+		t,
+		func(w io.Writer) *slog.Logger {
+			return slog.New(WrapHandler(slog.NewJSONHandler(w, nil)))
+		},
+		nil,
+	)
+}
+
+func TestHandlerWithAttrs(t *testing.T) {
+	testLogger(
+		t,
+		func(w io.Writer) *slog.Logger {
+			return slog.New(NewJSONHandler(w, nil)).
+				With("key1", "val1").
+				With(ext.LogKeyTraceID, "trace-id").
+				With(ext.LogKeySpanID, "span-id")
+		},
+		nil,
+	)
+}
+
+func TestHandlerWithGroup(t *testing.T) {
+	t.Run("simple", func(t *testing.T) {
+		testLogger(
+			t,
+			func(w io.Writer) *slog.Logger {
+				return slog.New(NewJSONHandler(w, nil)).
+					WithGroup("some-group").
+					With("key1", "val1")
+			},
+			func(t *testing.T, entry map[string]interface{}) {
+				assert.Equal(t, map[string]interface{}{
+					"key1": "val1",
+				}, entry["some-group"], "group entry not found")
+			},
+		)
+	})
+
+	t.Run("nested groups", func(t *testing.T) {
+		testLogger(
+			t,
+			func(w io.Writer) *slog.Logger {
+				return slog.New(NewJSONHandler(w, nil)).
+					With("key0", "val0").
+					WithGroup("group1").
+					With("key1", "val1").
+					WithGroup("group2").
+					With("key2", "val2").
+					With("key3", "val3")
+			},
+			func(t *testing.T, entry map[string]interface{}) {
+				assert.Equal(t, "val0", entry["key0"], "root level entry not found")
+				assert.Equal(t, map[string]interface{}{
+					"key1": "val1",
+					"group2": map[string]interface{}{
+						"key2": "val2",
+						"key3": "val3",
+					},
+				}, entry["group1"], "nested group entries not found")
+			},
+		)
 	})
 }