chunhao_receive/bootstrap/db.go

package bootstrap

import (
	"designs/config"
	"designs/global"
	"io"
	"log"
	"os"
	"time"

	"go.uber.org/zap"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func InitializeDB() *gorm.DB {
	// 根据驱动配置进行初始化
	switch config.Get("database.driver") {
	case "mysql":
		return initMySqlGorm()
	default:
		return initMySqlGorm()
	}
}

func UnInitializeDB(db *gorm.DB) {
	if db != nil {
		if sqldb, err := db.DB(); err == nil {
			sqldb.Close()
		} else {
			global.App.Log.Errorf("close database failed: %v", err)
		}
	}
}

func initMySqlGorm() *gorm.DB {

	UserName := config.Get("database.userName")
	Password := config.Get("database.password")
	Host := config.Get("database.host")
	Port := config.Get("database.port")
	Database := config.Get("database.database")
	Charset := config.Get("database.charset")

	dsn := UserName + ":" + Password + "@tcp(" + Host + ":" + Port + ")/" +
		Database + "?charset=" + Charset + "&parseTime=True&loc=Local"
	mysqlConfig := mysql.Config{
		DSN:                       dsn,   // DSN data source name
		DefaultStringSize:         191,   // string 类型字段的默认长度
		DisableDatetimePrecision:  true,  // 禁用 datetime 精度，MySQL 5.6 之前的数据库不支持
		DontSupportRenameIndex:    true,  // 重命名索引时采用删除并新建的方式，MySQL 5.7 之前的数据库和 MariaDB 不支持重命名索引
		DontSupportRenameColumn:   true,  // 用 `change` 重命名列，MySQL 8 之前的数据库和 MariaDB 不支持重命名列
		SkipInitializeWithVersion: false, // 根据版本自动配置
	}
	if db, err := gorm.Open(mysql.New(mysqlConfig), &gorm.Config{
		DisableForeignKeyConstraintWhenMigrating: true,            // 禁用自动创建外键约束
		Logger:                                   getGormLogger(), // 使用自定义 Logger
	}); err != nil {
		global.App.Log.Error("mysql connect failed, err:", zap.Any("err", err))
		return nil
	} else {
		sqlDB, _ := db.DB()
		sqlDB.SetMaxIdleConns(config.GetInt("database.maxIdleConns"))
		sqlDB.SetMaxOpenConns(config.GetInt("database.maxOpenConns"))
		initMySqlTables(db)
		return db
	}
}

func getGormLogger() logger.Interface {
	var logMode logger.LogLevel

	switch config.Get("database.logMode") {
	case "silent":
		logMode = logger.Silent
	case "error":
		logMode = logger.Error
	case "warn":
		logMode = logger.Warn
	case "info":
		logMode = logger.Info
	default:
		logMode = logger.Info
	}

	return logger.New(getGormLogWriter(), logger.Config{
		SlowThreshold:             2 * time.Second,                                 // 慢 SQL 阈值
		LogLevel:                  logMode,                                         // 日志级别
		IgnoreRecordNotFoundError: true,                                            // 忽略ErrRecordNotFound（记录未找到）错误
		Colorful:                  !config.GetBool("database.enableFileLogWriter"), // 禁用彩色打印
	})
}

// 自定义 gorm Writer
func getGormLogWriter() logger.Writer {
	var writer io.Writer

	// 是否启用日志文件
	if config.GetBool("database.enableFileLogWriter") {
		// 自定义 Writer
		// writer = &lumberjack.Logger{
		// 	Filename:   global.App.Config.Log.RootDir + "/" + global.App.Config.Database.LogFilename,
		// 	MaxSize:    global.App.Config.Log.MaxSize,
		// 	MaxBackups: global.App.Config.Log.MaxBackups,
		// 	MaxAge:     global.App.Config.Log.MaxAge,
		// 	Compress:   global.App.Config.Log.Compress,
		// }
		if global.App.LogWriter == nil {
			panic("log writer is nil")
		}
		writer = global.App.LogWriter
	} else {
		// 默认 Writer
		writer = os.Stdout
	}
	return log.New(writer, "\r\n", log.LstdFlags)
}

// 数据库表初始化
func initMySqlTables(db *gorm.DB) {
	err := db.AutoMigrate(
	// models.User{},
	// models.Media{},
	)
	if err != nil {
		global.App.Log.Error("migrate table failed", zap.Any("err", err))
		os.Exit(0)
	}
}