m7s/m7s.live/engine/v4@v4.14.7-0.20240211022911-f320e61803b9/db/db.go

// Copyright 2019 autocareai.com. All rights reserved.
// Use of this source code is governed by autocareai.com.

package db

import (
	"context"
	"fmt"
	"go.uber.org/zap"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/utils"
	"go.uber.org/zap/zapcore"
	"log"
	"os"
	"reflect"
	"time"
)

var (
	db *gorm.DB
)


// Colors
const (
	Reset       = "\033[0m"
	Red         = "\033[31m"
	Green       = "\033[32m"
	Yellow      = "\033[33m"
	Blue        = "\033[34m"
	Magenta     = "\033[35m"
	Cyan        = "\033[36m"
	White       = "\033[37m"
	BlueBold    = "\033[34;1m"
	MagentaBold = "\033[35;1m"
	RedBold     = "\033[31;1m"
	YellowBold  = "\033[33;1m"
)

// Writer log writer interface
type Writer interface {
	Printf(string, ...interface{})
}

type CustomLogger struct {
	Writer
	Config
	infoStr, warnStr, errStr            string
	traceStr, traceErrStr, traceWarnStr string
}

type Config struct {
	SlowThreshold             time.Duration
	Colorful                  bool
	IgnoreRecordNotFoundError bool
	LogLevel                  zapcore.Level
	TraceWithLevel            zapcore.Level
	Zap                       *zap.Logger
}

type callback struct{}

// 处理gorm v2.0 find 为空返回nil的问题
func (c callback) ParserError(db *gorm.DB) {
	if db.Error == nil && db.RowsAffected == 0 && (db.Statement.ReflectValue.Kind() == reflect.Slice || db.Statement.ReflectValue.Kind() == reflect.Struct) {
		db.Error = gorm.ErrRecordNotFound
	}
}

func RegisterCallback(db *gorm.DB) {
	_ = db.Callback().Query().After("gorm:query").Register("error_parser", callback{}.ParserError)
}

var gormLogLevelMap = map[logger.LogLevel]zapcore.Level{
	logger.Info:  zap.InfoLevel,
	logger.Warn:  zap.WarnLevel,
	logger.Error: zap.ErrorLevel,
}

// LogMode log mode
func (c CustomLogger) LogMode(level logger.LogLevel) logger.Interface {
	zapLevel, exist := gormLogLevelMap[level]
	if !exist {
		zapLevel = zap.DebugLevel
	}
	newlogger := c
	newlogger.LogLevel = zapLevel
	newlogger.TraceWithLevel = zapLevel
	return &newlogger
}

func (c CustomLogger) Info(ctx context.Context, msg string, data ...interface{}) {
	if c.Config.LogLevel <= zap.InfoLevel && c.Config.Zap != nil {
		c.Config.Zap.Sugar().Infof(msg, data...)
	}
}

func (c CustomLogger) Warn(ctx context.Context, msg string, data ...interface{}) {
	if c.Config.LogLevel <= zap.WarnLevel && c.Config.Zap != nil {
		c.Config.Zap.Sugar().Warnf(msg, data...)
	}
}

func (c CustomLogger) Error(ctx context.Context, msg string, data ...interface{}) {
	if c.Config.LogLevel <= zap.ErrorLevel && c.Config.Zap != nil {
		c.Config.Zap.Sugar().Errorf(msg, data...)
	}
}

func (c CustomLogger) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) {
	elapsed := time.Since(begin)
	sql, rows := fc()
	if c.Config.LogLevel <= zap.InfoLevel {
		c.Printf(c.traceStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, rows, sql)
	}
	if c.Config.Zap == nil {
		return
	}
	switch {
	case err != nil:
		c.Config.Zap.Error("sql: "+sql, zap.Float64("elapsed", elapsed.Seconds()), zap.Int64("rows", rows), zap.String("error", err.Error()))
	case c.Config.SlowThreshold != 0 && elapsed.Seconds() > c.Config.SlowThreshold.Seconds():
		c.Config.Zap.Warn("sql: "+sql, zap.Float64("elapsed", elapsed.Seconds()), zap.Int64("rows", rows), zap.Float64("threshold", c.Config.SlowThreshold.Seconds()))
	case c.Config.LogLevel == zap.DebugLevel:
		log := c.Config.Zap.Debug
		if c.Config.TraceWithLevel == zap.InfoLevel {
			log = c.Config.Zap.Info
		} else if c.Config.TraceWithLevel == zap.WarnLevel {
			log = c.Config.Zap.Warn
		} else if c.Config.TraceWithLevel == zap.ErrorLevel {
			log = c.Config.Zap.Error
		}
		log("sql: "+sql, zap.Float64("elapsed", elapsed.Seconds()), zap.Int64("rows", rows))
	}
}

// NewGormLogger 返回带 zap logger 的 GormLogger
func NewGormLogger(writer Writer, config Config) CustomLogger {
	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 = Green + "%s\n" + Reset + Green + "[info] " + Reset
		warnStr = BlueBold + "%s\n" + Reset + Magenta + "[warn] " + Reset
		errStr = Magenta + "%s\n" + Reset + Red + "[error] " + Reset
		traceStr = Green + "%s\n" + Reset + Yellow + "[%.3fms] " + BlueBold + "[rows:%v]" + Reset + " %s"
		traceWarnStr = Green + "%s " + Yellow + "%s\n" + Reset + RedBold + "[%.3fms] " + Yellow + "[rows:%v]" + Magenta + " %s" + Reset
		traceErrStr = RedBold + "%s " + MagentaBold + "%s\n" + Reset + Yellow + "[%.3fms] " + BlueBold + "[rows:%v]" + Reset + " %s"
	}
	return CustomLogger{
		Writer:       writer,
		Config:       config,
		infoStr:      infoStr,
		warnStr:      warnStr,
		errStr:       errStr,
		traceStr:     traceStr,
		traceWarnStr: traceWarnStr,
		traceErrStr:  traceErrStr,
	}
}


// Setup 建立连接
func Setup(user, passwd, addr, dbname, charset string, maxIdle, maxConn int, logMode bool,logger *zap.Logger) *gorm.DB {
	conf := Config{
		SlowThreshold:             time.Second,
		Colorful:                  true,
		IgnoreRecordNotFoundError: false,
		LogLevel:                  zap.ErrorLevel,
		TraceWithLevel:            zap.ErrorLevel,
		Zap:                       logger,
	}
	// 是否开启调试模式
	if logMode {
		conf.LogLevel = zap.DebugLevel
		conf.TraceWithLevel = zap.DebugLevel
		conf.SlowThreshold = time.Millisecond * 200
	}
	config := &gorm.Config{
		Logger: NewGormLogger(log.New(os.Stdout, "\r\n", log.LstdFlags), conf),
	}
	// 组装参数
	dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=%s&parseTime=True&loc=Local",
		user, passwd, addr, dbname, charset)
	// 打开新连接
	var err error
	db, err = gorm.Open(mysql.Open(dsn), config)
	if err != nil {
		log.Fatal("open mysql connection failed. err: ", err)
	}

	// 其他设置

	sqlDB, err := db.DB()
	if err != nil {
		log.Fatal("open mysql connection failed. err: ", err)
	}
	fmt.Println("数据库MaxIdle和MaxOpen：",maxIdle,maxConn)
	sqlDB.SetMaxIdleConns(maxIdle)
	sqlDB.SetMaxOpenConns(maxConn)
	RegisterCallback(db)
	return db
}

// DB 获取连接
func DB() *gorm.DB {
	return db
}

// Close 关闭连接
func Close() {
	if db != nil {
		sqlDB, _ := db.DB()
		_ = sqlDB.Close()
	}
}