feat: support logger with ctx. (fmt trace)

logger/logger_ctx.go

@@ -0,0 +1,171 @@
+package logger
+
+import (
+	"context"
+	"fmt"
+
+	"go.uber.org/zap"
+)
+
+type LoggerCtx interface {
+	Debugf(template string, args ...interface{})
+	Infof(template string, args ...interface{})
+	Warnf(template string, args ...interface{})
+	Errorf(template string, args ...interface{})
+	Debug(args ...interface{})
+	Info(args ...interface{})
+	Warn(args ...interface{})
+	Error(args ...interface{})
+
+	DebugfCtx(ctx context.Context, template string, args ...interface{})
+	InfofCtx(ctx context.Context, template string, args ...interface{})
+	WarnfCtx(ctx context.Context, template string, args ...interface{})
+	ErrorfCtx(ctx context.Context, template string, args ...interface{})
+	DebugCtx(ctx context.Context, template string, args ...interface{})
+	InfoCtx(ctx context.Context, template string, args ...interface{})
+	WarnCtx(ctx context.Context, template string, args ...interface{})
+	ErrorCtx(ctx context.Context, template string, args ...interface{})
+
+	Named(name string) LoggerCtx
+	With(fields ...zap.Field) LoggerCtx
+}
+
+var _ LoggerCtx = (*loggerCtx)(nil)
+
+type loggerCtx struct {
+	logger *zap.Logger
+	config
+}
+
+func NewLoggerCtx(opts ...CtxOption) *loggerCtx {
+	cfg := newDefaultConfig()
+	for _, o := range opts {
+		cfg = o.apply(cfg)
+	}
+
+	cfg.opts = append(cfg.opts, zap.AddCallerSkip(1))
+
+	return &loggerCtx{
+		logger: zap.New(cfg.core, cfg.opts...),
+		config: cfg,
+	}
+}
+
+func (l *loggerCtx) GetTrace(ctx context.Context) Trace {
+	var traceResult Trace
+
+	if l.parseTrace != nil {
+		traceResult = l.parseTrace(ctx)
+	}
+
+	return traceResult
+}
+
+func (l *loggerCtx) getExtraFields(ctx context.Context) (fields []zap.Field) {
+	if l.enabledTrace {
+		trace := l.GetTrace(ctx)
+		fields = append(fields, zap.String(l.traceKey, trace.TraceID), zap.String(l.spanKey, trace.SpanID))
+	}
+	return fields
+}
+
+func (l *loggerCtx) getMessage(template string, args []interface{}) string {
+	if len(args) == 0 {
+		return template
+	}
+
+	if template != "" {
+		return fmt.Sprintf(template, args...)
+	}
+
+	if len(args) == 1 {
+		if str, ok := args[0].(string); ok {
+			return str
+		}
+	}
+	return fmt.Sprint(args...)
+}
+
+func (l *loggerCtx) Named(name string) LoggerCtx {
+	return &loggerCtx{
+		logger: l.logger.Named(name),
+		config: l.config,
+	}
+}
+
+func (l *loggerCtx) Debugf(template string, args ...interface{}) {
+	l.logger.Debug(l.getMessage(template, args))
+}
+
+func (l *loggerCtx) DebugfCtx(ctx context.Context, template string, args ...interface{}) {
+	l.logger.Debug(l.getMessage(template, args), l.getExtraFields(ctx)...)
+}
+
+func (l *loggerCtx) Infof(template string, args ...interface{}) {
+	l.logger.Info(l.getMessage(template, args))
+}
+
+func (l *loggerCtx) Warnf(template string, args ...interface{}) {
+	l.logger.Warn(l.getMessage(template, args))
+}
+
+func (l *loggerCtx) Errorf(template string, args ...interface{}) {
+	l.logger.Error(l.getMessage(template, args))
+}
+
+func (l *loggerCtx) InfofCtx(ctx context.Context, template string, args ...interface{}) {
+	l.logger.Info(l.getMessage(template, args), l.getExtraFields(ctx)...)
+}
+
+func (l *loggerCtx) WarnfCtx(ctx context.Context, template string, args ...interface{}) {
+	l.logger.Warn(l.getMessage(template, args), l.getExtraFields(ctx)...)
+}
+
+func (l *loggerCtx) ErrorfCtx(ctx context.Context, template string, args ...interface{}) {
+	l.logger.Error(l.getMessage(template, args), l.getExtraFields(ctx)...)
+}
+
+func (l *loggerCtx) Debug(args ...interface{}) {
+	l.logger.Debug(l.getMessage("", args))
+}
+
+func (l *loggerCtx) DebugCtx(ctx context.Context, template string, args ...interface{}) {
+	l.logger.Debug(l.getMessage(template, args), l.getExtraFields(ctx)...)
+}
+
+func (l *loggerCtx) Error(args ...interface{}) {
+	l.logger.Error(l.getMessage("", args))
+}
+
+func (l *loggerCtx) ErrorCtx(ctx context.Context, template string, args ...interface{}) {
+	l.logger.Error(l.getMessage(template, args), l.getExtraFields(ctx)...)
+}
+
+func (l *loggerCtx) Info(args ...interface{}) {
+	l.logger.Info(l.getMessage("", args))
+}
+
+func (l *loggerCtx) InfoCtx(ctx context.Context, template string, args ...interface{}) {
+	l.logger.Info(l.getMessage(template, args), l.getExtraFields(ctx)...)
+}
+
+func (l *loggerCtx) Warn(args ...interface{}) {
+	l.logger.Warn(l.getMessage("", args))
+}
+
+func (l *loggerCtx) WarnCtx(ctx context.Context, template string, args ...interface{}) {
+	l.logger.Warn(l.getMessage("", args), l.getExtraFields(ctx)...)
+}
+
+func (l *loggerCtx) With(fields ...zap.Field) LoggerCtx {
+	if len(fields) == 0 {
+		return l
+	}
+	temp := l.clone()
+	temp.logger = l.logger.With(fields...)
+	return temp
+}
+
+func (l *loggerCtx) clone() *loggerCtx {
+	return &loggerCtx{config: l.config}
+}

logger/logger_test.go

@@ -7,13 +7,17 @@ package logger
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
+	"io"
+	"log"
 	"os"
 	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap/zapcore"
 )
 
 func BenchmarkMuitiLogs(b *testing.B) {
@@ -419,6 +423,96 @@ func TestRotateOnDevNull(t *testing.T) {
 	}
 }
 
+type BufferSync struct {
+	io.ReadWriter
+}
+
+func (b *BufferSync) Sync() error {
+	return nil
+}
+
+func TestTrace(t *testing.T) {
+	tests := []struct {
+		name            string
+		enabled         bool
+		traceLoggerName string
+		spanLoggerName  string
+		message         string
+		extractTrace    ExtractTrace
+		expectedOutput  string
+	}{
+		{
+			name:            "parse trace",
+			enabled:         true,
+			traceLoggerName: "traceID",
+			spanLoggerName:  "span_id",
+			message:         "test",
+			extractTrace: func(ctx context.Context) Trace {
+				return Trace{
+					SpanID:  "2",
+					TraceID: "1",
+				}
+			},
+			expectedOutput: `{"level":"DEBUG","message":"test","traceID":"1","span_id":"2"}` + "\n",
+		},
+		{
+			name:           "not trace",
+			enabled:        false,
+			message:        "test",
+			extractTrace:   nil,
+			expectedOutput: `{"level":"DEBUG","message":"test"}` + "\n",
+		},
+		{
+			name:            "not parse",
+			enabled:         true,
+			traceLoggerName: "traceID",
+			spanLoggerName:  "span_id",
+			message:         "test",
+			extractTrace:    nil,
+			expectedOutput:  `{"level":"DEBUG","message":"test","traceID":"","span_id":""}` + "\n",
+		},
+	}
+
+	buf := bytes.NewBufferString("")
+	bsync := &BufferSync{buf}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			ast := assert.New(t)
+			opt := make([]CtxOption, 0)
+			opt = append(opt,
+				WithTraceKey(test.traceLoggerName, test.spanLoggerName),
+				WithParseTrace(test.extractTrace),
+			)
+			if test.enabled {
+				opt = append(opt, EnableTrace())
+			}
+
+			opt = append(opt, WithZapCore(
+				zapcore.NewCore(
+					zapcore.NewJSONEncoder(zapcore.EncoderConfig{
+						MessageKey:  "message",
+						LevelKey:    "level",
+						EncodeLevel: zapcore.CapitalLevelEncoder,
+					}),
+					bsync,
+					zapcore.DebugLevel,
+				)))
+
+			logger := NewLoggerCtx(opt...)
+
+			ctx := context.Background()
+
+			logger.DebugfCtx(ctx, test.message)
+			re, err := io.ReadAll(bsync)
+			if err != nil {
+				log.Fatal(err)
+			}
+			ast.Equal(test.expectedOutput, string(re))
+		})
+	}
+}
+
 //nolint:forbidigo
 func showLog(f string) {
 	logdata, err := os.ReadFile(f)

logger/opt.go

@@ -0,0 +1,83 @@
+package logger
+
+import (
+	"context"
+	"os"
+
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+)
+
+const (
+	defaultTraceLoggerName = "trace_id"
+	defaultSpanLoggerName  = "span_id"
+)
+
+type CtxOption interface {
+	apply(config) config
+}
+
+type optionFunc func(config) config
+
+func (fn optionFunc) apply(cfg config) config {
+	return fn(cfg)
+}
+
+type Trace struct {
+	TraceID string
+	SpanID  string
+}
+
+type ExtractTrace func(ctx context.Context) Trace
+
+type config struct {
+	enabledTrace bool
+	traceKey     string
+	spanKey      string
+	parseTrace   ExtractTrace
+	core         zapcore.Core
+	opts         []zap.Option
+}
+
+func newDefaultConfig() config {
+	return config{
+		traceKey: defaultSpanLoggerName,
+		core:     zapcore.NewCore(zapcore.NewJSONEncoder(zapcore.EncoderConfig{MessageKey: "message"}), os.Stdout, zap.DebugLevel),
+	}
+}
+
+func EnableTrace() CtxOption {
+	return optionFunc(func(cfg config) config {
+		cfg.enabledTrace = true
+		return cfg
+	})
+}
+
+func WithZapCore(core zapcore.Core) CtxOption {
+	return optionFunc(func(cfg config) config {
+		cfg.core = core
+		return cfg
+	})
+}
+
+func WithZapOpts(opts ...zap.Option) CtxOption {
+	return optionFunc(func(cfg config) config {
+		cfg.opts = opts
+		return cfg
+	})
+}
+
+func WithTraceKey(tn string, sn string) CtxOption {
+	return optionFunc(func(cfg config) config {
+		cfg.traceKey = tn
+		cfg.spanKey = sn
+		return cfg
+	})
+}
+
+func WithParseTrace(f ExtractTrace) CtxOption {
+	return optionFunc(func(cfg config) config {
+		cfg.parseTrace = f
+		return cfg
+	})
+}