feat(logger): add WithFields/FieldsFromContext methods for context

Author: ulya-sidorina

CHANGELOG.md

@@ -1,3 +1,7 @@
+## v3.116.4
+* Added helper methods `WithFields` and `FieldsFromContext` for working with structured logging fields via context.
+  These methods allow adding custom fields to the context, which are later extracted by the logger.
+
 ## v3.116.3
 * Default grpc message size adjusted to the server's defaults
 

log/context.go

@@ -4,9 +4,12 @@ import (
 	"context"
 )
 
+type FieldsType map[string]string
+
 type (
-	ctxLevelKey struct{}
-	ctxNamesKey struct{}
+	ctxLevelKey  struct{}
+	ctxNamesKey  struct{}
+	ctxFieldsKey struct{}
 )
 
 func WithLevel(ctx context.Context, lvl Level) context.Context {
@@ -36,6 +39,29 @@ func NamesFromContext(ctx context.Context) []string {
 	return v[:len(v):len(v)] // prevent re
 }
 
+func WithFields(ctx context.Context, fields FieldsType) context.Context {
+	existing := FieldsFromContext(ctx)
+	merged := make(FieldsType, len(existing)+len(fields))
+
+	for k, v := range existing {
+		merged[k] = v
+	}
+
+	for k, v := range fields {
+		merged[k] = v
+	}
+
+	return context.WithValue(ctx, ctxFieldsKey{}, merged)
+}
+
+func FieldsFromContext(ctx context.Context) FieldsType {
+	if fields, _ := ctx.Value(ctxFieldsKey{}).(FieldsType); fields != nil {
+		return fields
+	}
+
+	return FieldsType{}
+}
+
 func with(ctx context.Context, lvl Level, names ...string) context.Context {
 	return WithLevel(WithNames(ctx, names...), lvl)
 }

log/context_test.go

@@ -90,3 +90,37 @@ func TestWithNamesRaceRegression(t *testing.T) {
 		}
 	})
 }
+
+func TestFieldsFromContext(t *testing.T) {
+	for _, tt := range []struct {
+		ctx    context.Context //nolint:containedctx
+		fields FieldsType
+	}{
+		{
+			ctx:    context.Background(),
+			fields: FieldsType{},
+		},
+		{
+			ctx:    WithFields(context.Background(), FieldsType{"a": "1", "b": "1"}),
+			fields: FieldsType{"a": "1", "b": "1"},
+		},
+		{
+			ctx: WithFields(
+				WithFields(context.Background(), FieldsType{"a": "1", "b": "1"}),
+				FieldsType{"a": "2", "b": "2"},
+			),
+			fields: FieldsType{"a": "2", "b": "2"},
+		},
+		{
+			ctx: WithFields(
+				WithFields(context.Background(), FieldsType{"a": "1", "b": "1"}),
+				FieldsType{"a": "3"},
+			),
+			fields: FieldsType{"a": "3", "b": "1"},
+		},
+	} {
+		t.Run("", func(t *testing.T) {
+			require.Equal(t, tt.fields, FieldsFromContext(tt.ctx))
+		})
+	}
+}

log/logger.go

@@ -90,9 +90,23 @@ func (l *defaultLogger) Log(ctx context.Context, msg string, fields ...Field) {
 		return
 	}
 
+	additionalFields := FieldsFromContext(ctx)
+	var merged []Field
+	if len(additionalFields) > 0 {
+		merged = make([]Field, 0, len(additionalFields)+len(fields))
+
+		for k, v := range additionalFields {
+			merged = append(merged, String(k, v))
+		}
+
+		merged = append(merged, fields...)
+	} else {
+		merged = fields
+	}
+
 	_, _ = io.WriteString(l.w, l.format(
 		NamesFromContext(ctx),
-		l.appendFields(msg, fields...),
+		l.appendFields(msg, merged...),
 		lvl,
 	)+"\n")
 }