mycommon/mymysql/mysql.go

package mymysql

import (
	"database/sql"
	"fmt"
	"git.makemake.in/kzkzzzz/mycommon/myconf"
	driverMysql "github.com/go-sql-driver/mysql"
	"github.com/google/uuid"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"log"
	"os"
	"sync"
	"time"
)

const (
	DefaultInstance = "mysql"
)

type MysqlDb struct {
	*gorm.DB
	SqlDB       *sql.DB
	gormConfig  *gorm.Config
	disablePing bool
}

type Conf struct {
	Dsn              string
	MaxOpenConn      int    // 最大连接数
	MaxIdleConn      int    // 最大空闲连接数
	MaxIdleTime      string // 空闲时间
	MaxLifeTime      string // 连接最大有效时间
	Debug            bool
	LogSqlSlowTimeMs int
	LogDisableColor  bool
}

var (
	instanceMap = &sync.Map{}
)

func GetDb(name ...string) *MysqlDb {
	var instanceName string
	if len(name) > 0 {
		instanceName = name[0]
	} else {
		instanceName = DefaultInstance
	}

	v, ok := instanceMap.Load(instanceName)
	if !ok {
		panic(fmt.Errorf("mysql instance [%s] not init", instanceName))
	}

	return v.(*MysqlDb)
}

// InitDb 初始化全局默认db
func InitDb(config *myconf.Config, opts ...Opt) {
	client, err := NewDb(DefaultInstance, config, opts...)
	if err != nil {
		panic(err)
	}
	instanceMap.Store(DefaultInstance, client)
}

// InitDbInstance 初始化全局的db
func InitDbInstance(instanceName string, config *myconf.Config, opts ...Opt) {
	client, err := NewDb(instanceName, config, opts...)
	if err != nil {
		panic(err)
	}
	instanceMap.Store(instanceName, client)
}

func NewDb(instanceName string, config *myconf.Config, opts ...Opt) (*MysqlDb, error) {
	cf := &Conf{Debug: true}
	err := config.UnmarshalKey(instanceName, &cf)
	if err != nil {
		return nil, err
	}
	db, err := NewDbFromConf(cf, opts...)
	if err != nil {
		return nil, err
	}

	instanceMap.Store(instanceName, db)
	return db, nil
}

func NewDbFromConf(cf *Conf, opts ...Opt) (*MysqlDb, error) {
	parseDsn, err := driverMysql.ParseDSN(cf.Dsn)
	if err != nil {
		return nil, fmt.Errorf("mysql parse dsn error: %s", err)
	}

	db := &MysqlDb{}
	for _, opt := range opts {
		opt(db)
	}

	if db.gormConfig == nil {
		db.gormConfig = &gorm.Config{
			SkipDefaultTransaction: true,
		}

		lCfg := logger.Config{
			SlowThreshold:             200 * time.Millisecond,
			LogLevel:                  logger.Warn,
			IgnoreRecordNotFoundError: false,
			Colorful:                  true,
		}

		if cf.LogSqlSlowTimeMs > 0 {
			lCfg.SlowThreshold = time.Duration(cf.LogSqlSlowTimeMs) * time.Millisecond
		}

		if cf.LogDisableColor {
			lCfg.Colorful = false
		}

		l := newGormLogger(lCfg)

		if cf.Debug {
			db.gormConfig.Logger = l.LogMode(logger.Info)
		} else {
			db.gormConfig.Logger = l
		}
	}

	gormDB, err := gorm.Open(mysql.Open(cf.Dsn), db.gormConfig)
	if err != nil {
		return nil, err
	}

	sqlDB, err := gormDB.DB()
	if err != nil {
		return nil, err
	}

	if db.disablePing == false {
		err = sqlDB.Ping()
		if err != nil {
			return nil, err
		}
	}

	if cf.MaxOpenConn <= 0 {
		cf.MaxOpenConn = 1024
	}

	if cf.MaxIdleConn <= 0 {
		// 默认最大空闲数等于最大连接数
		cf.MaxIdleConn = cf.MaxOpenConn
	}

	if cf.MaxIdleTime == "" {
		cf.MaxIdleTime = "10m"
	}

	sqlDB.SetMaxOpenConns(cf.MaxOpenConn)
	sqlDB.SetMaxIdleConns(cf.MaxIdleConn)

	if dv, err := time.ParseDuration(cf.MaxIdleTime); err != nil {
		return nil, fmt.Errorf("parse MaxIdleTime err: %s", err)
	} else {
		sqlDB.SetConnMaxIdleTime(dv)
	}

	// max life time默认暂不设置, 使用idle time控制即可
	if cf.MaxLifeTime != "" {
		if dv, err := time.ParseDuration(cf.MaxLifeTime); err != nil {
			return nil, fmt.Errorf("parse MaxLifeTime err: %s", err)
		} else {
			sqlDB.SetConnMaxLifetime(dv)
		}
	}

	db.DB = gormDB

	db.SqlDB = sqlDB
	instanceMap.Store(uuid.New().String(), db)

	log.Printf("connect db success [addr:%s - db:%s]", parseDsn.Addr, parseDsn.DBName)

	return db, nil
}

func CloseAll() {
	instanceMap.Range(func(k, v any) bool {
		db, err := (v.(*MysqlDb)).DB.DB()
		if err != nil {
			db.Close()
		}
		return true
	})
}

func newGormLogger(cfg logger.Config) logger.Interface {
	return logger.New(log.New(os.Stdout, "\r\n", log.LstdFlags), cfg)
}