feat(logging/zap): log extrakeys (hertz-contrib#26)

Signed-off-by: rogerogers <rogers@rogerogers.com>

logging/zap/logger.go

@@ -66,6 +66,20 @@ func NewLogger(opts ...Option) *Logger {
 	}
 }
 
+// GetExtraKeys get extraKeys from logger config
+func (l *Logger) GetExtraKeys() []ExtraKey {
+	return l.config.extraKeys
+}
+
+// PutExtraKeys add extraKeys after init
+func (l *Logger) PutExtraKeys(keys ...ExtraKey) {
+	for _, k := range keys {
+		if !inArray(k, l.config.extraKeys) {
+			l.config.extraKeys = append(l.config.extraKeys, k)
+		}
+	}
+}
+
 func (l *Logger) Log(level klog.Level, kvs ...interface{}) {
 	logger := l.With()
 	switch level {
@@ -123,6 +137,16 @@ func (l *Logger) CtxLogf(level klog.Level, ctx context.Context, format string, k
 		sl = l.With()
 	}
 
+	if len(l.config.extraKeys) > 0 {
+		for _, k := range l.config.extraKeys {
+			if l.config.extraKeyAsStr {
+				sl = sl.With(string(k), ctx.Value(string(k)))
+			} else {
+				sl = sl.With(string(k), ctx.Value(k))
+			}
+		}
+	}
+
 	switch level {
 	case klog.LevelDebug, klog.LevelTrace:
 		zlevel = zap.DebugLevel

logging/zap/logger_test.go

@@ -17,6 +17,7 @@ package zap
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"strings"
 	"testing"
 
@@ -273,3 +274,66 @@ func TestCtxKVLogger(t *testing.T) {
 		buf.Reset()
 	}
 }
+
+// TestWithExtraKeys test WithExtraKeys option
+func TestWithExtraKeys(t *testing.T) {
+	buf := new(bytes.Buffer)
+
+	log := NewLogger(WithExtraKeys([]ExtraKey{"requestId"}))
+	log.SetOutput(buf)
+
+	ctx := context.WithValue(context.Background(), ExtraKey("requestId"), "123")
+
+	log.CtxInfof(ctx, "%s log", "extra")
+
+	var logStructMap map[string]interface{}
+
+	err := json.Unmarshal(buf.Bytes(), &logStructMap)
+
+	assert.Nil(t, err)
+
+	value, ok := logStructMap["requestId"]
+
+	assert.True(t, ok)
+	assert.Equal(t, value, "123")
+}
+
+func TestPutExtraKeys(t *testing.T) {
+	logger := NewLogger(WithExtraKeys([]ExtraKey{"abc"}))
+
+	assert.Contains(t, logger.GetExtraKeys(), ExtraKey("abc"))
+	assert.NotContains(t, logger.GetExtraKeys(), ExtraKey("def"))
+
+	logger.PutExtraKeys("def")
+	assert.Contains(t, logger.GetExtraKeys(), ExtraKey("def"))
+}
+
+func TestExtraKeyAsStr(t *testing.T) {
+	buf := new(bytes.Buffer)
+	const v = "value"
+
+	logger := NewLogger(WithExtraKeys([]ExtraKey{"abc"}))
+
+	logger.SetOutput(buf)
+
+	ctx1 := context.TODO()
+	ctx1 = context.WithValue(ctx1, "key1", v) //nolint:staticcheck
+	logger.CtxErrorf(ctx1, "%s", "error")
+
+	assert.NotContains(t, buf.String(), v)
+
+	buf.Reset()
+
+	strLogger := NewLogger(WithExtraKeys([]ExtraKey{"abc"}), WithExtraKeyAsStr())
+
+	strLogger.SetOutput(buf)
+
+	ctx2 := context.TODO()
+	ctx2 = context.WithValue(ctx2, "key2", v) //nolint:staticcheck
+
+	strLogger.CtxErrorf(ctx2, "key2", v)
+
+	assert.Contains(t, buf.String(), v)
+
+	buf.Reset()
+}

logging/zap/option.go

@@ -21,6 +21,8 @@ import (
 	"go.uber.org/zap/zapcore"
 )
 
+type ExtraKey string
+
 type Option interface {
 	apply(cfg *config)
 }
@@ -43,9 +45,11 @@ type traceConfig struct {
 }
 
 type config struct {
-	coreConfig  coreConfig
-	zapOpts     []zap.Option
-	traceConfig *traceConfig
+	extraKeys     []ExtraKey
+	coreConfig    coreConfig
+	zapOpts       []zap.Option
+	traceConfig   *traceConfig
+	extraKeyAsStr bool
 }
 
 // defaultCoreConfig default zapcore config: json encoder, atomic level, stdout write syncer
@@ -73,7 +77,8 @@ func defaultConfig() *config {
 			recordStackTraceInSpan: true,
 			errorSpanLevel:         zapcore.ErrorLevel,
 		},
-		zapOpts: []zap.Option{},
+		zapOpts:       []zap.Option{},
+		extraKeyAsStr: false,
 	}
 }
 
@@ -118,3 +123,25 @@ func WithRecordStackTraceInSpan(recordStackTraceInSpan bool) Option {
 		cfg.traceConfig.recordStackTraceInSpan = recordStackTraceInSpan
 	})
 }
+
+// WithExtraKeys allow you log extra values from context
+func WithExtraKeys(keys []ExtraKey) Option {
+	return option(func(cfg *config) {
+		for _, k := range keys {
+			if !inArray(k, cfg.extraKeys) {
+				cfg.extraKeys = append(cfg.extraKeys, k)
+			}
+		}
+	})
+}
+
+// WithExtraKeyAsStr convert extraKey to a string type when retrieving value from context
+// Not recommended for use, only for compatibility with certain situations
+//
+// For more information, refer to the documentation at
+// `https://pkg.go.dev/context#WithValue`
+func WithExtraKeyAsStr() Option {
+	return option(func(cfg *config) {
+		cfg.extraKeyAsStr = true
+	})
+}

logging/zap/utils.go

@@ -47,3 +47,13 @@ func OtelSeverityText(lv zapcore.Level) string {
 	}
 	return s
 }
+
+// inArray check if a string in a slice
+func inArray(key ExtraKey, arr []ExtraKey) bool {
+	for _, k := range arr {
+		if k == key {
+			return true
+		}
+	}
+	return false
+}