From ddbf7a04844eebd9b20b95cc15bba1f398532084 Mon Sep 17 00:00:00 2001
From: HF <crazytaxii666@gmail.com>
Date: Tue, 21 May 2024 06:45:27 -0400
Subject: [PATCH] Trace all SQL statements and output with log

---
 api/server/middleware/log.go        | 10 ++--
 api/server/middleware/middleware.go |  2 +-
 cmd/app/options/options.go          |  9 ++--
 pkg/db/logger.go                    | 73 +++++++++++++++++++++++++++++
 4 files changed, 86 insertions(+), 8 deletions(-)
 create mode 100644 pkg/db/logger.go

diff --git a/api/server/middleware/log.go b/api/server/middleware/log.go
index 206cafde..e9226113 100644
--- a/api/server/middleware/log.go
+++ b/api/server/middleware/log.go
@@ -23,6 +23,8 @@ import (
 	"github.com/gin-contrib/requestid"
 	"github.com/gin-gonic/gin"
 	klog "github.com/sirupsen/logrus"
+
+	"github.com/caoyingjunz/pixiu/pkg/db"
 )
 
 const (
@@ -31,9 +33,10 @@ const (
 	FailMsg    = "FAIL"
 )
 
-func LoggerToFile() gin.HandlerFunc {
+func Logger() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		startTime := time.Now()
+		c.Set(db.SQLContextKey, new(db.SQLs)) // set SQL context key
 
 		// 处理请求操作
 		c.Next()
@@ -45,8 +48,9 @@ func LoggerToFile() gin.HandlerFunc {
 			"status_code": c.Writer.Status(),
 			"latency":     fmt.Sprintf("%dµs", time.Since(startTime).Microseconds()),
 			"client_ip":   c.ClientIP(),
-			// TODO
-			// "sqls": []string{},
+		}
+		if sqls := db.GetSQLs(c); len(sqls) > 0 {
+			fields["sqls"] = sqls
 		}
 
 		if errs := c.Errors; len(errs) > 0 {
diff --git a/api/server/middleware/middleware.go b/api/server/middleware/middleware.go
index 10256db6..87e1e9bc 100644
--- a/api/server/middleware/middleware.go
+++ b/api/server/middleware/middleware.go
@@ -53,7 +53,7 @@ func InstallMiddlewares(o *options.Options) {
 			return util.GenerateRequestID()
 		})),
 		Cors(),
-		LoggerToFile(),
+		Logger(),
 		UserRateLimiter(),
 		Limiter(),
 		Authentication(o.ComponentConfig.Default),
diff --git a/cmd/app/options/options.go b/cmd/app/options/options.go
index 5aa9d4cc..fd247f28 100644
--- a/cmd/app/options/options.go
+++ b/cmd/app/options/options.go
@@ -19,6 +19,7 @@ package options
 import (
 	"fmt"
 	"os"
+	"time"
 
 	pixiuConfig "github.com/caoyingjunz/pixiulib/config"
 	"github.com/gin-gonic/gin"
@@ -40,6 +41,8 @@ const (
 	defaultTokenKey   = "pixiu"
 	defaultConfigFile = "/etc/pixiu/config.yaml"
 	defaultLogFormat  = config.LogFormatJson
+
+	defaultSlowSQLDuration = 1 * time.Second
 )
 
 // Options has all the params needed to run a pixiu
@@ -131,11 +134,9 @@ func (o *Options) registerDatabase() error {
 		sqlConfig.Port,
 		sqlConfig.Name)
 
-	opt := &gorm.Config{}
-	if o.ComponentConfig.Default.Mode == "debug" {
-		opt.Logger = logger.Default.LogMode(logger.Info)
+	opt := &gorm.Config{
+		Logger: db.NewLogger(logger.Info, defaultSlowSQLDuration),
 	}
-
 	DB, err := gorm.Open(mysql.Open(dsn), opt)
 	if err != nil {
 		return err
diff --git a/pkg/db/logger.go b/pkg/db/logger.go
new file mode 100644
index 00000000..2f180399
--- /dev/null
+++ b/pkg/db/logger.go
@@ -0,0 +1,73 @@
+/*
+Copyright 2024 The Pixiu Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package db
+
+import (
+	"context"
+	"time"
+
+	"gorm.io/gorm/logger"
+)
+
+type (
+	SQLs []string
+
+	DBLogger struct {
+		logger.LogLevel
+		SlowThreshold time.Duration // slow SQL queries
+	}
+)
+
+const SQLContextKey = "sqls"
+
+func NewLogger(level logger.LogLevel, slowThreshold time.Duration) *DBLogger {
+	return &DBLogger{
+		LogLevel:      level,
+		SlowThreshold: slowThreshold,
+	}
+}
+
+func (l *DBLogger) LogMode(level logger.LogLevel) logger.Interface {
+	l.LogLevel = level
+	return l
+}
+
+func (l *DBLogger) Info(ctx context.Context, msg string, data ...interface{}) {}
+
+func (l *DBLogger) Warn(ctx context.Context, msg string, data ...interface{}) {}
+
+func (l *DBLogger) Error(ctx context.Context, msg string, data ...interface{}) {}
+
+func (l *DBLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
+	if l.LogLevel <= logger.Silent {
+		return
+	}
+
+	sql, _ := fc()
+	if v := ctx.Value(SQLContextKey); v != nil {
+		sqls := v.(*SQLs)
+		*sqls = append(*sqls, sql)
+	}
+}
+
+// GetSQLs returns all the SQL statements executed in the current context.
+func GetSQLs(ctx context.Context) SQLs {
+	if v := ctx.Value(SQLContextKey); v != nil {
+		return *v.(*SQLs)
+	}
+	return SQLs{}
+}