97 lines
2.5 KiB
Go
97 lines
2.5 KiB
Go
package utils_gorm
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"gorm.io/gorm"
|
||
"gorm.io/gorm/logger"
|
||
"regexp"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
type CustomLogger struct {
|
||
gormLogger logger.Interface
|
||
db *gorm.DB
|
||
}
|
||
|
||
func NewCustomLogger(db *gorm.DB) *CustomLogger {
|
||
return &CustomLogger{
|
||
gormLogger: logger.Default.LogMode(logger.Info),
|
||
db: db,
|
||
}
|
||
}
|
||
|
||
func (l *CustomLogger) LogMode(level logger.LogLevel) logger.Interface {
|
||
newlogger := *l
|
||
newlogger.gormLogger = l.gormLogger.LogMode(level)
|
||
return &newlogger
|
||
}
|
||
|
||
func (l *CustomLogger) Info(ctx context.Context, msg string, data ...interface{}) {
|
||
l.gormLogger.Info(ctx, msg, data...)
|
||
}
|
||
|
||
func (l *CustomLogger) Warn(ctx context.Context, msg string, data ...interface{}) {
|
||
l.gormLogger.Warn(ctx, msg, data...)
|
||
}
|
||
|
||
func (l *CustomLogger) Error(ctx context.Context, msg string, data ...interface{}) {
|
||
l.gormLogger.Error(ctx, msg, data...)
|
||
}
|
||
|
||
func (l *CustomLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
|
||
elapsed := time.Since(begin)
|
||
sql, _ := fc()
|
||
l.gormLogger.Trace(ctx, begin, fc, err)
|
||
operation := extractOperation(sql)
|
||
tableName := extractTableName(sql)
|
||
fmt.Println(tableName)
|
||
//// 将SQL语句保存到数据库
|
||
if operation == 0 || tableName == "sql_log" {
|
||
return
|
||
}
|
||
//go l.db.Model(&SqlLog{}).Create(&SqlLog{
|
||
// OperatorID: 1,
|
||
// OperatorName: "test",
|
||
// SqlInfo: sql,
|
||
// TableNames: tableName,
|
||
// Type: operation,
|
||
//})
|
||
|
||
// 如果有需要，也可以根据执行时间（elapsed）等条件过滤或处理日志记录
|
||
if elapsed > time.Second {
|
||
//l.gormLogger.Warn(ctx, "Slow SQL (> 1s): %s", sql)
|
||
}
|
||
}
|
||
|
||
// extractTableName extracts the table name from a SQL query, supporting quoted table names.
|
||
func extractTableName(sql string) string {
|
||
// 使用非捕获组匹配多种SQL操作关键词
|
||
re := regexp.MustCompile(`(?i)\b(?:from|update|into|delete\s+from)\b\s+[\` + "`" + `"]?(\w+)[\` + "`" + `"]?`)
|
||
match := re.FindStringSubmatch(sql)
|
||
|
||
// 检查是否匹配成功
|
||
if len(match) > 1 {
|
||
return match[1]
|
||
}
|
||
|
||
return ""
|
||
}
|
||
|
||
// extractOperation extracts the operation type from a SQL query.
|
||
func extractOperation(sql string) int32 {
|
||
sql = strings.TrimSpace(strings.ToLower(sql))
|
||
var operation int32
|
||
if strings.HasPrefix(sql, "select") {
|
||
operation = 0
|
||
} else if strings.HasPrefix(sql, "insert") {
|
||
operation = 1
|
||
} else if strings.HasPrefix(sql, "update") {
|
||
operation = 3
|
||
} else if strings.HasPrefix(sql, "delete") {
|
||
operation = 2
|
||
}
|
||
return operation
|
||
}
|