mycommon/mymysql/mysql.go

package mymysql

import (
	"fmt"
	"git.makemake.in/kzkzzzz/mycommon/mylog"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	gormLogger "gorm.io/gorm/logger"
	"time"
)

const DefaultKey = "default"

var (
	DefaultConfig = &Config{
		Dsn:         "root:root@tcp(127.0.0.1:3306)/?loc=Local&charset=utf8mb4&parseTime=true",
		MaxOpenConn: 32,
		MaxIdleConn: 8,
		MaxLifeTime: "4h",
		MaxIdleTime: "15m",
		Debug:       true,
		GormLogger:  gormLogger.Default.LogMode(gormLogger.Info),
	}

	instanceMap = make(map[string]*gorm.DB)
)

type (
	Config struct {
		Dsn         string
		MaxOpenConn int
		MaxIdleConn int
		MaxIdleTime string
		MaxLifeTime string
		Debug       bool
		GormLogger  gormLogger.Interface
	}
)

func DB(key ...string) *gorm.DB {
	var key0 string

	if len(key) > 0 {
		key0 = key[0]
	} else {
		key0 = DefaultKey
	}

	instance, ok := instanceMap[key0]
	if !ok {
		panic(fmt.Errorf("mysql %s not config", key0))
	}
	return instance
}

func InitDefault(config *Config) error {
	return Init(DefaultKey, config)
}

func Init(key string, config *Config) error {
	db, err := New(config)
	if err != nil {
		return err
	}
	instanceMap[key] = db
	return nil
}

func New(config *Config) (*gorm.DB, error) {
	var (
		maxLifeTime, _ = time.ParseDuration(DefaultConfig.MaxLifeTime)
		maxIdleTime, _ = time.ParseDuration(DefaultConfig.MaxIdleTime)
		logger         gormLogger.Interface
	)

	if config.MaxOpenConn <= 0 {
		config.MaxOpenConn = DefaultConfig.MaxOpenConn
	}

	if config.MaxIdleConn <= 0 {
		config.MaxIdleConn = DefaultConfig.MaxIdleConn
	}

	if config.MaxLifeTime != "" {
		t, err := time.ParseDuration(config.MaxLifeTime)
		if err != nil {
			return nil, fmt.Errorf("parse MaxLifeTime err: %s\n", err)

		}
		maxLifeTime = t
	}

	if config.MaxIdleTime != "" {
		t, err := time.ParseDuration(config.MaxIdleTime)
		if err != nil {
			return nil, fmt.Errorf("parse MaxIdleTime err: %s\n", err)
		}
		maxIdleTime = t
	}

	if config.GormLogger == nil {
		level := gormLogger.Warn
		if config.Debug {
			level = gormLogger.Info
		}
		logger = DefaultGormLogger(level)
	}

	db, err := gorm.Open(mysql.Open(config.Dsn), &gorm.Config{
		SkipDefaultTransaction: true,
		Logger:                 logger,
	})

	if err != nil {
		return nil, fmt.Errorf("connect mysql err: %s", err)
	}
	sqlDb, _ := db.DB()

	sqlDb.SetMaxOpenConns(config.MaxOpenConn)
	sqlDb.SetMaxIdleConns(config.MaxIdleConn)
	sqlDb.SetConnMaxLifetime(maxLifeTime)
	sqlDb.SetConnMaxIdleTime(maxIdleTime)

	return db, nil
}

func DefaultGormLogger(level gormLogger.LogLevel) gormLogger.Interface {
	return gormLogger.New(mylog.NewLogger("gorm", mylog.DefaultConfig), gormLogger.Config{
		SlowThreshold:             time.Second * 2,
		Colorful:                  true,
		IgnoreRecordNotFoundError: false,
		ParameterizedQueries:      false,
		LogLevel:                  level,
	})
}

func NewGormLogger(writer gormLogger.Writer, gormLoggerConfig gormLogger.Config) gormLogger.Interface {
	return gormLogger.New(writer, gormLoggerConfig)
}

func CloseDB(key string) {
	db, _ := DB(key).DB()
	db.Close()
}

func CloseAllDB() {
	for _, v := range instanceMap {
		db, _ := v.DB()
		db.Close()
	}
}