cadmin/server/config/database.go

package config

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

	mysql2 "github.com/go-sql-driver/mysql"
	"github.com/sirupsen/logrus"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"gorm.io/plugin/dbresolver"
)

var (
	DB       *gorm.DB
	LDB      *gorm.DB
	GDB1     *gorm.DB
	GDBGroup map[int]*gorm.DB
	//AdminDB  *gorm.DB
)

func InitDatabase() {
	initDefaultDB()
	initGraveDB()
	initLoginDB()
	//initAdminDB()
}

func initDefaultDB() {
	connString := os.Getenv("MYSQL_DSN")
	newLogger := logger.New(
		log.New(os.Stdout, "\r\n", log.LstdFlags), // io writer
		logger.Config{
			SlowThreshold:             time.Second, // Slow SQL threshold
			LogLevel:                  logger.Info, // Log level(这里记得根据需求改一下)
			IgnoreRecordNotFoundError: true,        // Ignore ErrRecordNotFound error for logger
			Colorful:                  false,       // Disable color
		},
	)

	db, err := gorm.Open(mysql.Open(connString), &gorm.Config{
		Logger:                 newLogger,
		SkipDefaultTransaction: true,
	})

	if connString == "" || err != nil {
		logrus.Fatalf("mysql lost: %v", err)
	}
	sqlDB, err := db.DB()
	if err != nil {
		logrus.Fatalf("mysql lost: %v", err)
	}

	sqlDB.SetMaxIdleConns(10)
	sqlDB.SetMaxOpenConns(20)
	sqlDB.SetConnMaxLifetime(2 * time.Hour)
	DB = db
}

//func initAdminDB() {
//	connString := os.Getenv("ADMIN_MYSQL_DSN")
//	newLogger := logger.New(
//		log.New(os.Stdout, "\r\n", log.LstdFlags), // io writer
//		logger.Config{
//			SlowThreshold:             time.Second, // Slow SQL threshold
//			LogLevel:                  logger.Info, // Log level(这里记得根据需求改一下)
//			IgnoreRecordNotFoundError: true,        // Ignore ErrRecordNotFound error for logger
//			Colorful:                  false,       // Disable color
//		},
//	)
//
//	db, err := gorm.Open(mysql.Open(connString), &gorm.Config{
//		Logger:                 newLogger,
//		SkipDefaultTransaction: true,
//	})
//
//	if connString == "" || err != nil {
//		logrus.Fatalf("mysql lost: %v", err)
//	}
//	sqlDB, err := db.DB()
//	if err != nil {
//		logrus.Fatalf("mysql lost: %v", err)
//	}
//
//	sqlDB.SetMaxIdleConns(10)
//	sqlDB.SetMaxOpenConns(20)
//	sqlDB.SetConnMaxLifetime(2 * time.Hour)
//	AdminDB = db
//}

func initGraveDB() {
	GDBGroup = make(map[int]*gorm.DB)

	GDB1 = initSingleGraveDB(os.Getenv("MYSQL_GRAVE_DSN"), os.Getenv("MYSQL_GRAVE_SLAVE_DSN"), 1)
	GDBGroup[1] = GDB1

	// 游戏服数据库组，新加服务器时往后续就行
	// 当连续3各数据库都不存在时，停止遍历

	failNum := 0
	for i := 2; i < 10000; i += 1 {
		conn := os.Getenv(fmt.Sprintf("MYSQL_GRAVE%v_DSN", i))
		slave := os.Getenv(fmt.Sprintf("MYSQL_GRAVE%v_SLAVE_DSN", i))
		if conn == "" || slave == "" {
			failNum += 1
			if failNum > 3 {
				return
			}
			//logrus.Warnf("initSingleGraveDB empty index[%v] present, skipped..", index)
			continue
		}
		failNum = 0
		initSingleGraveDB(conn, slave, i)
	}
}

func initSingleGraveDB(conn, slave string, index int) *gorm.DB {
	//logrus.Debugf("initSingleGraveDB index:%v, conn:%v, slave:%v", index, conn, slave)
	newLogger := logger.New(
		log.New(os.Stdout, "\r\n", log.LstdFlags), // io writer
		logger.Config{
			SlowThreshold:             time.Second, // Slow SQL threshold
			LogLevel:                  logger.Info, // Log level(这里记得根据需求改一下)
			IgnoreRecordNotFoundError: true,        // Ignore ErrRecordNotFound error for logger
			Colorful:                  false,       // Disable color
		},
	)

	db1, err := gorm.Open(mysql.Open(conn), &gorm.Config{
		Logger:                 newLogger,
		SkipDefaultTransaction: true,
	})
	if conn == "" || err != nil {
		logrus.Fatalf("mysql lost: %v", err)
	}

	// 读写分离
	err = db1.Use(dbresolver.Register(dbresolver.Config{
		Sources:  []gorm.Dialector{mysql.Open(conn)},
		Replicas: []gorm.Dialector{mysql.Open(slave)},
	}))
	if slave == "" || err != nil {
		logrus.Fatalf("mysql dbresolver slaveString：%v, err: %v", slave, err)
	}

	sqlDB, err := db1.DB()
	if err != nil {
		logrus.Fatalf("mysql lost: %v", err)
	}

	sqlDB.SetMaxIdleConns(20)
	sqlDB.SetMaxOpenConns(50)
	sqlDB.SetConnMaxLifetime(2 * time.Hour)

	if index > 0 {
		GDBGroup[index] = db1
	}
	return db1
}

func initLoginDB() {
	connString := os.Getenv("MYSQL_LOGIN_DSN")
	newLogger := logger.New(
		log.New(os.Stdout, "\r\n", log.LstdFlags), // io writer
		logger.Config{
			SlowThreshold:             time.Second, // Slow SQL threshold
			LogLevel:                  logger.Info, // Log level(这里记得根据需求改一下)
			IgnoreRecordNotFoundError: true,        // Ignore ErrRecordNotFound error for logger
			Colorful:                  false,       // Disable color
		},
	)

	db, err := gorm.Open(mysql.Open(connString), &gorm.Config{
		Logger:                 newLogger,
		SkipDefaultTransaction: true,
	})
	if connString == "" || err != nil {
		logrus.Fatalf("mysql lost: %v", err)
	}

	// 读写分离
	slaveString := os.Getenv("MYSQL_LOGIN_SLAVE_DSN")
	err = db.Use(dbresolver.Register(dbresolver.Config{
		Sources:  []gorm.Dialector{mysql.Open(connString)},
		Replicas: []gorm.Dialector{mysql.Open(slaveString)},
	}))
	if slaveString == "" || err != nil {
		logrus.Fatalf("mysql dbresolver slaveString：%v, err: %v", slaveString, err)
	}

	sqlDB, err := db.DB()
	if err != nil {
		logrus.Fatalf("mysql lost: %v", err)
	}

	sqlDB.SetMaxIdleConns(20)
	sqlDB.SetMaxOpenConns(50)
	sqlDB.SetConnMaxLifetime(2 * time.Hour)
	LDB = db
}

// GetDBName 获取数据库库名
func GetDBName(conn string) string {
	dsnConf, _ := mysql2.ParseDSN(conn)
	return dsnConf.DBName
}

func GetDBId(DB *gorm.DB) int {
	for k, db := range GDBGroup {
		if db == DB {
			return k
		}
	}
	return 0
}