From a70f77464ca02d17db74baf7a973e07342acad70 Mon Sep 17 00:00:00 2001
From: douxu <dou921247973@outlook.com>
Date: Mon, 23 Jun 2025 16:00:48 +0800
Subject: [PATCH] refactor(gorm-logger): 1. add gorm logger in gorm config
 2.use faced func in gorm logger

---
 database/postgres_init.go |  4 ++-
 logger/gorm_logger.go     | 61 +++++++++++++++++++++++++++++++++++++++
 main.go                   |  8 +++--
 3 files changed, 69 insertions(+), 4 deletions(-)
 create mode 100644 logger/gorm_logger.go

diff --git a/database/postgres_init.go b/database/postgres_init.go
index d4edcc8..a80d9fc 100644
--- a/database/postgres_init.go
+++ b/database/postgres_init.go
@@ -6,6 +6,8 @@ import (
 	"sync"
 	"time"
 
+	"modelRT/logger"
+
 	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
 )
@@ -36,7 +38,7 @@ func InitPostgresDBInstance(ctx context.Context, PostgresDBURI string) *gorm.DB
 func initPostgresDBClient(ctx context.Context, PostgresDBURI string) *gorm.DB {
 	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
-	db, err := gorm.Open(postgres.Open(PostgresDBURI), &gorm.Config{})
+	db, err := gorm.Open(postgres.Open(PostgresDBURI), &gorm.Config{Logger: logger.NewGormLogger()})
 	if err != nil {
 		panic(err)
 	}
diff --git a/logger/gorm_logger.go b/logger/gorm_logger.go
new file mode 100644
index 0000000..511d677
--- /dev/null
+++ b/logger/gorm_logger.go
@@ -0,0 +1,61 @@
+// Package logger define log struct of modelRT project
+package logger
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"gorm.io/gorm"
+	gormLogger "gorm.io/gorm/logger"
+)
+
+// GormLogger define struct for implementing gormLogger.Interface
+type GormLogger struct {
+	SlowThreshold time.Duration
+}
+
+// NewGormLogger define func for init GormLogger
+func NewGormLogger() *GormLogger {
+	return &GormLogger{
+		SlowThreshold: 500 * time.Millisecond,
+	}
+}
+
+// LogMode define func for implementing gormLogger.Interface
+func (l *GormLogger) LogMode(_ gormLogger.LogLevel) gormLogger.Interface {
+	return &GormLogger{}
+}
+
+// Info define func for implementing gormLogger.Interface
+func (l *GormLogger) Info(ctx context.Context, msg string, data ...any) {
+	Info(ctx, msg, "data", data)
+}
+
+// Warn define func for implementing gormLogger.Interface
+func (l *GormLogger) Warn(ctx context.Context, msg string, data ...any) {
+	Warn(ctx, msg, "data", data)
+}
+
+// Error define func for implementing gormLogger.Interface
+func (l *GormLogger) Error(ctx context.Context, msg string, data ...any) {
+	Error(ctx, msg, "data", data)
+}
+
+// Trace define func for implementing gormLogger.Interface
+func (l *GormLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
+	// get SQL running time
+	duration := time.Since(begin).Milliseconds()
+	// get gorm exec sql and rows affected
+	sql, rows := fc()
+	// gorm error judgment
+	if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
+		Error(ctx, "SQL ERROR", "sql", sql, "rows", rows, "dur(ms)", duration)
+	}
+	// slow query judgment
+	if duration > l.SlowThreshold.Milliseconds() {
+		Warn(ctx, "SQL SLOW", "sql", sql, "rows", rows, "dur(ms)", duration)
+	} else {
+		Debug(ctx, "SQL DEBUG", "sql", sql, "rows", rows, "dur(ms)", duration)
+	}
+}
diff --git a/main.go b/main.go
index ccc53a2..a058a44 100644
--- a/main.go
+++ b/main.go
@@ -17,13 +17,12 @@ import (
 	"modelRT/middleware"
 	"modelRT/pool"
 
-	swaggerFiles "github.com/swaggo/files"
-	ginSwagger "github.com/swaggo/gin-swagger"
-
 	realtimedata "modelRT/real-time-data"
 
 	"github.com/gin-gonic/gin"
 	"github.com/panjf2000/ants/v2"
+	swaggerFiles "github.com/swaggo/files"
+	ginSwagger "github.com/swaggo/gin-swagger"
 	"gorm.io/gorm"
 )
 
@@ -50,6 +49,9 @@ func main() {
 	flag.Parse()
 	ctx := context.TODO()
 
+	// init logger
+	logger.InitLoggerInstance(modelRTConfig.LoggerConfig)
+
 	modelRTConfig = config.ReadAndInitConfig(*modelRTConfigDir, *modelRTConfigName, *modelRTConfigType)
 	// init postgresDBClient
 	postgresDBClient = database.InitPostgresDBInstance(ctx, modelRTConfig.PostgresDBURI)