pkg.daxia.dev/db/init_db.go

package db

import (
	"fmt"
	"log"
	"os"
	"time"

	"gorm.io/driver/mysql"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/schema"

	"github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"gorm.io/gorm"
)

type Database struct {
	Self *gorm.DB
}

// 单例
var DB *Database

func InitTest() {
	viper.Set("gormlog", viper.GetBool("gormlog"))
	viper.Set("db.username", viper.GetString("test.db.username"))
	viper.Set("db.password", viper.GetString("test.db.password"))
	viper.Set("db.addr", viper.GetString("test.db.addr"))
	viper.Set("db.name", viper.GetString("test.db.name"))

	DB = &Database{
		Self: getDB(),
	}
}

func Init() {
	DB = &Database{
		Self: getDB(),
	}
}

func openDB(username, password, addr, name string) *gorm.DB {
	config := fmt.Sprintf(
		"%s:%s@tcp(%s)/%s?charset=utf8mb4&parseTime=%t&loc=%s&timeout=%ds",
		username,
		password,
		addr,
		name,
		true,
		//"Asia%2FShanghai",  // 必须是 url.QueryEscape 的
		"Local",
		10,
	)

	logLevel := logger.Warn

	newLogger := logger.New(log.New(os.Stdout, "\r\n", log.LstdFlags), logger.Config{
		SlowThreshold:             200 * time.Millisecond,
		LogLevel:                  logLevel,
		IgnoreRecordNotFoundError: false,
		Colorful:                  true,
	}).LogMode(logger.Error)

	// if viper.GetBool("gormlog") {
	// 	newLogger = logger.Default.LogMode(logger.Info)
	// }

	if viper.GetBool("gormlog") {
		newLogger = logger.Default.LogMode(logger.Info)
	}

	logrus.Info("config:", config)

	db, err := gorm.Open(mysql.Open(config), &gorm.Config{
		Logger: newLogger,
		NamingStrategy: schema.NamingStrategy{
			TablePrefix:   viper.GetString("db.tableprefix"),
			SingularTable: viper.GetBool("db.singulartable"),
		},
	})

	if err != nil {
		logrus.Fatalf("数据库连接失败. 数据库名字: %s. 错误信息: %s", name, err)
	} else {
		logrus.Infof("数据库链接成功, 数据库名字: %s", name)
	}

	setupDB(db)
	return db
}

func setupDB(db *gorm.DB) {
	//db.LogMode(viper.GetBool("gormlog"))
	// 用于设置最大打开的连接数，默认值为0表示不限制.设置最大的连接数，可以避免并发太高导致连接mysql出现too many connections的错误。
	//db.DB().SetMaxOpenConns(20000)
	// 用于设置闲置的连接数.设置闲置的连接数则当开启的一个连接使用完成后可以放在池里等候下一次使用。
	//db.DB().SetMaxIdleConns(60)
	//db.SingularTable(true)@@
	sqlDB, err := db.DB()
	if err != nil {
		panic(err)
	}

	sqlDB.SetMaxIdleConns(100)

}

func InitDB() *gorm.DB {
	return openDB(
		viper.GetString("db.username"),
		viper.GetString("db.password"),
		viper.GetString("db.addr"),
		viper.GetString("db.name"),
	)
}

func getDB() *gorm.DB {
	if DB != nil {
		return DB.Self
	}

	return InitDB()
}

func GetDB() *gorm.DB {
	return DB.Self
}