自定义gorm的logger

server/config.yaml

@@ -60,6 +60,7 @@ mysql:
   max-idle-conns: 10
   max-open-conns: 100
   log-mode: false
+  log-zap: false
 
 # local configuration
 local:

server/config/gorm.go

@@ -9,4 +9,5 @@ type Mysql struct {
 	MaxIdleConns int    `mapstructure:"max-idle-conns" json:"maxIdleConns" yaml:"max-idle-conns"`
 	MaxOpenConns int    `mapstructure:"max-open-conns" json:"maxOpenConns" yaml:"max-open-conns"`
 	LogMode      bool   `mapstructure:"log-mode" json:"logMode" yaml:"log-mode"`
+	LogZap       bool   `mapstructure:"log-zap" json:"logZap" yaml:"log-zap"`
 }

server/initialize/gorm.go

@@ -73,6 +73,12 @@ func GormMysql() *gorm.DB {
 
 // gormConfig 根据配置决定是否开启日志
 func gormConfig(mod bool) *gorm.Config {
+	if global.GVA_CONFIG.Mysql.LogZap {
+		return &gorm.Config{
+			Logger:                                   Default.LogMode(logger.Info),
+			DisableForeignKeyConstraintWhenMigrating: true,
+		}
+	}
 	if mod {
 		return &gorm.Config{
 			Logger:                                   logger.Default.LogMode(logger.Info),

server/initialize/logger.go

@@ -0,0 +1,146 @@
+package initialize
+
+import (
+	"context"
+	"fmt"
+	"gin-vue-admin/global"
+	"go.uber.org/zap"
+	"gorm.io/gorm/logger"
+	"gorm.io/gorm/utils"
+	"io/ioutil"
+	"log"
+	"os"
+	"time"
+)
+
+var (
+	Discard = New(log.New(ioutil.Discard, "", log.LstdFlags), GormConfig{})
+	Default = New(log.New(os.Stdout, "\r\n", log.LstdFlags), GormConfig{
+		SlowThreshold: 200 * time.Millisecond,
+		LogLevel:      logger.Warn,
+		Colorful:      true,
+	})
+	Recorder = traceRecorder{Interface: Default, BeginAt: time.Now()}
+)
+
+type traceRecorder struct {
+	logger.Interface
+	BeginAt      time.Time
+	SQL          string
+	RowsAffected int64
+	Err          error
+}
+
+func New(writer Writer, config GormConfig) logger.Interface {
+	var (
+		infoStr      = "%s\n[info] "
+		warnStr      = "%s\n[warn] "
+		errStr       = "%s\n[error] "
+		traceStr     = "%s\n[%.3fms] [rows:%v] %s"
+		traceWarnStr = "%s %s\n[%.3fms] [rows:%v] %s"
+		traceErrStr  = "%s %s\n[%.3fms] [rows:%v] %s"
+	)
+
+	if config.Colorful {
+		infoStr = logger.Green + "%s\n" + logger.Reset + logger.Green + "[info] " + logger.Reset
+		warnStr = logger.BlueBold + "%s\n" + logger.Reset + logger.Magenta + "[warn] " + logger.Reset
+		errStr = logger.Magenta + "%s\n" + logger.Reset + logger.Red + "[error] " + logger.Reset
+		traceStr = logger.Green + "%s\n" + logger.Reset + logger.Yellow + "[%.3fms] " + logger.BlueBold + "[rows:%v]" + logger.Reset + " %s"
+		traceWarnStr = logger.Green + "%s " + logger.Yellow + "%s\n" + logger.Reset + logger.RedBold + "[%.3fms] " + logger.Yellow + "[rows:%v]" + logger.Magenta + " %s" + logger.Reset
+		traceErrStr = logger.RedBold + "%s " + logger.MagentaBold + "%s\n" + logger.Reset + logger.Yellow + "[%.3fms] " + logger.BlueBold + "[rows:%v]" + logger.Reset + " %s"
+	}
+
+	return &GormLogger{
+		Writer:       writer,
+		GormConfig:   config,
+		infoStr:      infoStr,
+		warnStr:      warnStr,
+		errStr:       errStr,
+		traceStr:     traceStr,
+		traceWarnStr: traceWarnStr,
+		traceErrStr:  traceErrStr,
+	}
+}
+
+// Writer log writer interface
+type Writer interface {
+	Printf(string, ...interface{})
+}
+
+type GormConfig struct {
+	SlowThreshold time.Duration
+	Colorful      bool
+	LogLevel      logger.LogLevel
+}
+
+type GormLogger struct {
+	Writer
+	GormConfig
+	infoStr, warnStr, errStr            string
+	traceStr, traceErrStr, traceWarnStr string
+}
+
+func (g *GormLogger) LogMode(level logger.LogLevel) logger.Interface {
+	newLogger := *g
+	newLogger.LogLevel = level
+	return &newLogger
+}
+
+func (g *GormLogger) Info(ctx context.Context, message string, data ...interface{}) {
+	if g.LogLevel >= logger.Info {
+		g.Printf(g.infoStr+message, append([]interface{}{utils.FileWithLineNum()}, data...)...)
+	}
+}
+
+func (g *GormLogger) Warn(ctx context.Context, message string, data ...interface{}) {
+	if g.LogLevel >= logger.Warn {
+		g.Printf(g.warnStr+message, append([]interface{}{utils.FileWithLineNum()}, data...)...)
+	}
+}
+
+func (g *GormLogger) Error(ctx context.Context, message string, data ...interface{}) {
+	if g.LogLevel >= logger.Error {
+		g.Printf(g.errStr+message, append([]interface{}{utils.FileWithLineNum()}, data...)...)
+	}
+}
+
+func (g *GormLogger) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) {
+	if g.LogLevel > 0 {
+		elapsed := time.Since(begin)
+		switch {
+		case err != nil && g.LogLevel >= logger.Error:
+			sql, rows := fc()
+			if rows == -1 {
+				g.Printf(g.traceErrStr, utils.FileWithLineNum(), err, float64(elapsed.Nanoseconds())/1e6, "-", sql)
+			} else {
+				g.Printf(g.traceErrStr, utils.FileWithLineNum(), err, float64(elapsed.Nanoseconds())/1e6, rows, sql)
+			}
+		case elapsed > g.SlowThreshold && g.SlowThreshold != 0 && g.LogLevel >= logger.Warn:
+			sql, rows := fc()
+			slowLog := fmt.Sprintf("SLOW SQL >= %v", g.SlowThreshold)
+			if rows == -1 {
+				g.Printf(g.traceWarnStr, utils.FileWithLineNum(), slowLog, float64(elapsed.Nanoseconds())/1e6, "-", sql)
+			} else {
+				g.Printf(g.traceWarnStr, utils.FileWithLineNum(), slowLog, float64(elapsed.Nanoseconds())/1e6, rows, sql)
+			}
+		case g.LogLevel >= logger.Info:
+			sql, rows := fc()
+			if rows == -1 {
+				g.Printf(g.traceStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, "-", sql)
+			} else {
+				g.Printf(g.traceStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, rows, sql)
+			}
+		}
+	}
+}
+
+func (g *GormLogger) Printf(message string, data ...interface{}) {
+	global.GVA_LOG.Info(
+		"gorm",
+		zap.String("type", "sql"),
+		zap.Any("src", data[0]),
+		zap.Any("duration", data[1]),
+		zap.Any("rows", data[2]),
+		zap.Any("sql", data[3]),
+	)
+}