From d41516d05d3ddd29b501355073154b8f9fd47b88 Mon Sep 17 00:00:00 2001
From: Yun <995116474@qq.com>
Date: Wed, 3 Jul 2024 16:58:14 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96tace=5Fid?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 loggerx.go            |  5 +++++
 middleware/traceId.go | 15 +++++++++++++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/loggerx.go b/loggerx.go
index 8fa57db..cdebd9c 100644
--- a/loggerx.go
+++ b/loggerx.go
@@ -87,6 +87,11 @@ func (l *Logger) Channel(ch string) (r *Logger) {
 	return &rr
 }
 
+// 获取TraceField的字段
+func (l *Logger) GetTraceField() string {
+	return l.option.traceField
+}
+
 // 实现io.Writer接口
 func (l *Logger) Write(b []byte) (n int, err error) {
 	return l.write("info", b)
diff --git a/middleware/traceId.go b/middleware/traceId.go
index deec9e3..5936010 100644
--- a/middleware/traceId.go
+++ b/middleware/traceId.go
@@ -5,10 +5,11 @@ import (
 
 	"github.com/gin-gonic/gin"
 	uuid "github.com/satori/go.uuid"
+	"github.com/yuninks/loggerx"
 )
 
 // 设置普通的traceId
-func SetTraceId(ctx context.Context, traceKey string) context.Context {
+func SetTraceIdByKey(ctx context.Context, traceKey string) context.Context {
 	if traceKey == "" {
 		traceKey = "trace_id"
 	}
@@ -20,8 +21,13 @@ func SetTraceId(ctx context.Context, traceKey string) context.Context {
 	return ctx
 }
 
+// 设置logger的traceId
+func SetTraceId(ctx context.Context, logger *loggerx.Logger) context.Context {
+	return SetTraceIdByKey(ctx, logger.GetTraceField())
+}
+
 // 设置Gin的traceId
-func SetGinTraceId(traceKey string) gin.HandlerFunc {
+func SetGinTraceIdByKey(traceKey string) gin.HandlerFunc {
 
 	if traceKey == "" {
 		traceKey = "trace_id"
@@ -35,3 +41,8 @@ func SetGinTraceId(traceKey string) gin.HandlerFunc {
 		ctx.Set(traceKey, traceId)
 	}
 }
+
+// 设置Gin的traceId
+func SetGinTraceIdByLogger(logger *loggerx.Logger) gin.HandlerFunc {
+	return SetGinTraceIdByKey(logger.GetTraceField())
+}