添加 WithContext 方法

Author: hwholiday

hlog/zap.go

@@ -3,12 +3,13 @@ package hlog
 import (
 	"context"
 	"fmt"
-	"github.com/natefinch/lumberjack"
-	"go.uber.org/zap"
-	"go.uber.org/zap/zapcore"
 	"os"
 	"strings"
 	"sync"
+
+	"github.com/natefinch/lumberjack"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
 )
 
 var (
@@ -52,6 +53,14 @@ func (l *Logger) GetCtx(ctx context.Context) *zap.Logger {
 	return l.Logger
 }
 
+func (l *Logger) WithContext(ctx context.Context) *zap.Logger {
+	log, ok := ctx.Value(l.opts.CtxKey).(*zap.Logger)
+	if ok {
+		return log
+	}
+	return l.Logger
+}
+
 func (l *Logger) AddCtx(ctx context.Context, field ...zap.Field) (context.Context, *zap.Logger) {
 	log := l.With(field...)
 	ctx = context.WithValue(ctx, l.opts.CtxKey, log)