🐛 fix log middleware, which can't record query args right.

kamiertop · Mar 19, 2024 · c881393 · c881393
1 parent 6329584
commit c881393
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 11 deletions.
diff --git a/pkg/middleware/zap_log.go b/pkg/middleware/zap_log.go
@@ -3,6 +3,7 @@ package middleware
 import (
 	"context"
 	"net/http"
+	"net/url"
 	"time"
 
 	"github.com/cloudwego/hertz/pkg/app"
@@ -22,13 +23,16 @@ func Logger() app.HandlerFunc {
 		logMsg := []zap.Field{
 			zap.String("route", c.FullPath()),
 			zap.String("ip", c.RemoteAddr().String()),
-			zap.String("agent", utils.BytesToString(c.UserAgent())),
 		}
-		requestArg(c, &logMsg)
+		logMsg = requestArg(c, logMsg)
 		s := time.Now()
 		c.Next(ctx)
 		cost := time.Since(s)
-		logMsg = append(logMsg, zap.String("cost", cost.String()), zap.Int("status", c.Response.StatusCode()))
+
+		logMsg = append(logMsg,
+			zap.String("cost", cost.String()),
+			zap.Int("status", c.Response.StatusCode()),
+			zap.String("agent", utils.BytesToString(c.UserAgent())))
 		switch c.Response.StatusCode() {
 		case http.StatusOK:
 			if cost.Seconds() > _execTimeout {
@@ -38,28 +42,29 @@ func Logger() app.HandlerFunc {
 			}
 		default:
 			if e := c.Errors.Last(); e != nil {
-				logMsg = append(logMsg, zap.Error(e))
+				logMsg = append(logMsg, zap.Error(e.Err), zap.String("err_msg", e.Meta.(string)))
 			}
 			log.Error(string(c.Request.URI().Path()), logMsg...)
 		}
 	}
 }
 
-func requestArg(c *app.RequestContext, logMsg *[]zap.Field) {
+func requestArg(c *app.RequestContext, logMsg []zap.Field) []zap.Field {
 	switch string(c.Method()) {
 	case http.MethodGet:
-		if v := c.QueryArgs(); v.Len() != 0 {
-			*logMsg = append(*logMsg, zap.String("query_rgs", utils.BytesToString(v.QueryString())))
+		query, err := url.QueryUnescape(c.QueryArgs().String())
+		if err == nil && query != "" {
+			logMsg = append(logMsg, zap.String("query_rgs", query))
 		}
 		if p := c.Params; len(p) != 0 {
-			*logMsg = append(*logMsg, zap.Any("params", c.Params))
+			logMsg = append(logMsg, zap.Any("params", c.Params))
 		}
 	case http.MethodPost:
 		if body := c.Request.Body(); body != nil {
 			if utils.BytesToString(c.GetHeader(consts.HeaderContentType)) == consts.MIMEApplicationJSON {
 				var b map[string]any
 				if err := json.Unmarshal(body, &b); err == nil {
-					*logMsg = append(*logMsg, zap.Any("body", b))
+					logMsg = append(logMsg, zap.Any("body", b))
 				}
 			}
 		}
@@ -69,4 +74,6 @@ func requestArg(c *app.RequestContext, logMsg *[]zap.Field) {
 	default:
 
 	}
+
+	return logMsg
 }
diff --git a/pkg/storage/mongo.go b/pkg/storage/mongo.go
@@ -9,9 +9,11 @@ import (
 	"go.mongodb.org/mongo-driver/mongo/readpref"
 )
 
+var client *mongo.Client
+
 // InitMongo init mongo driver.
-func InitMongo() error {
-	client, err := mongo.Connect(context.Background(), &options.ClientOptions{})
+func InitMongo() (err error) {
+	client, err = mongo.Connect(context.Background(), &options.ClientOptions{})
 	if err != nil {
 		return fmt.Errorf("connect mongo error: %w", err)
 	}
@@ -22,3 +24,18 @@ func InitMongo() error {
 
 	return nil
 }
+
+// Client return mongo client.
+func Client() *mongo.Client {
+	return client
+}
+
+// Database return database by name.
+func Database(name string) *mongo.Database {
+	return client.Database(name)
+}
+
+// CloseMongo disconnect mongo.
+func CloseMongo(ctx context.Context) error {
+	return client.Disconnect(ctx)
+}