feat(logger): add WithFields method to context

Author: ulya-sidorina

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
+		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))
+		})
+	}
+}