eta_api/global/initDb.go

package global

import (
	"database/sql/driver"
	_ "dm"
	dm "dmgorm2"
	"eta/eta_api/utils"
	"fmt"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/schema"
	"io"
	"log"
	"os"
	"strings"
	"time"
)

var (
	//需要钩子函数执行的数据库表
	tableMap = map[string]int{
		"base_from_mysteel_chemical_data": 1,
		"future_good_edb_info":            1,
	}
)

type LocalTime time.Time

// InitDb
// @Description: 数据库初始化
func InitDb() {
	dbMap := make(map[string]*gorm.DB)

	//开启日志
	logWriter := io.MultiWriter(utils.Binlog) //binlog日志，记录到文件中去
	if utils.RunMode == `dev` {               // 测试环境，默认输出在控制台，不需要的话，可以注释if里面下面的这行
		logWriter = io.MultiWriter(utils.Binlog, os.Stdout)
	}
	newLogger := logger.New(log.New(logWriter, "\r\n", log.LstdFlags), logger.Config{
		SlowThreshold:             200 * time.Millisecond, //慢sql ：200ms
		LogLevel:                  logger.Info,            //记录的日志类型，info代表所有信息都记录
		IgnoreRecordNotFoundError: true,                   //是否忽略找不到数据错误信息（只是日志记录记录成err还是普通的输出的区别，并不影响业务代码中的：找不到数据行error）
		Colorful:                  false,                  //是否颜色输出
	})

	// 默认库
	connectDb(utils.MYSQL_URL, utils.DbNameMaster, newLogger, dbMap, true)

	// 报告库
	connectDb(utils.MYSQL_URL_RDDP, utils.DbNameReport, newLogger, dbMap, false)
	// 手工数据库
	connectDb(utils.MYSQL_URL_EDB, utils.DbNameManualIndex, newLogger, dbMap, false)
	// 指标库
	connectDb(utils.MYSQL_URL_DATA, utils.DbNameIndex, newLogger, dbMap, false)
	// 钢联库
	connectDb(utils.MYSQL_URL_GL, utils.DbNameGL, newLogger, dbMap, false)

	if utils.MYSQL_AI_URL != "" {
		// AI库
		connectDb(utils.MYSQL_AI_URL, utils.DbNameAI, newLogger, dbMap, false)
	}

	// 用户主库
	if utils.MYSQL_WEEKLY_URL != `` && (utils.BusinessCode == utils.BusinessCodeRelease || utils.BusinessCode == utils.BusinessCodeSandbox || utils.BusinessCode == utils.BusinessCodeDebug) {
		connectDb(utils.MYSQL_WEEKLY_URL, utils.DbNameWeekly, newLogger, dbMap, false)
	}
	for _, db := range dbMap {
		_ = db.Callback().Query().Before("gorm:query").Register("before_query", func(tx *gorm.DB) {
			orgSql := tx.Statement.SQL.String()
			if _, ok := tableMap[tx.Statement.Table]; ok && strings.Contains(strings.ToLower(orgSql), "select") {
				newSql := utils.ReplaceDriverKeywords(utils.DbDriverName, orgSql)
				tx.Statement.SQL.Reset()
				tx.Statement.SQL.WriteString(newSql)
			}
		})
		_ = db.Callback().Raw().Before("gorm:raw").Register("before_raw", func(tx *gorm.DB) {
			orgSql := tx.Statement.SQL.String()
			if _, ok := tableMap[tx.Statement.Table]; ok &&
				(strings.Contains(strings.ToLower(orgSql), "delete") ||
					strings.Contains(strings.ToLower(orgSql), "update") ||
					strings.Contains(strings.ToLower(orgSql), "insert")) {
				newSql := utils.ReplaceDriverKeywords(utils.DbDriverName, orgSql)
				tx.Statement.SQL.Reset()
				tx.Statement.SQL.WriteString(newSql)
			}
		})
	}
	//全局赋值数据库链接
	DbMap = dbMap

}

// connectDb
// @Description: 达梦数据库连接
// @param dsn
// @param aliasName
// @param newLogger
// @param dbMap
// @param isDefault
func connectDb(dsn, aliasName string, newLogger logger.Interface, dbMap map[string]*gorm.DB, isDefault bool) {
	//fmt.Println("dsn:", dsn, "  ==  ;aliasName:", aliasName)
	if dsn == `` {
		return
	}

	var dialector gorm.Dialector
	switch utils.DbDriverName {
	case utils.DbDriverByMysql:
		if !strings.Contains(dsn, `parseTime`) {
			dsn += `&parseTime=true`
		}
		dialector = mysql.Open(dsn)
	case utils.DbDriverByDm:
		dialector = dm.Open(dsn)
	default:
		panic(fmt.Errorf("数据库 链接异常，错误的数据库类型，数据库:%s", utils.DbDriverName))
	}
	db, err := gorm.Open(dialector, &gorm.Config{
		Logger: newLogger,
		NamingStrategy: schema.NamingStrategy{
			SingularTable: true, // 表示使用单数表名，启用该选项后，GORM 将不会对表名进行复数化处理
		},
	})
	if err != nil {
		//global.LOG.Errorf("mysql 启动异常，数据库:default;Err：", err)
		panic(fmt.Errorf("数据库 链接异常，数据库:%s;Err:%s", aliasName, err))
	}
	//创建连接池
	sqlDB, err := db.DB()
	if err != nil {
		//global.LOG.Errorf("达梦 创建连接池失败，数据库:default;Err：", err)
		panic(fmt.Errorf("数据库 创建连接池失败，数据库:%s;Err:%s", aliasName, err))
	}

	dbMap[aliasName] = db

	//默认数据库连接
	if isDefault {
		DEFAULT_DB = db
	}

	// SetMaxIdleConns 设置空闲连接池中连接的最大数量
	sqlDB.SetMaxIdleConns(50)

	// SetMaxOpenConns 设置打开数据库连接的最大数量。
	sqlDB.SetMaxOpenConns(100)

	// SetConnMaxLifetime 设置了连接可复用的最大时间。
	sqlDB.SetConnMaxLifetime(10 * time.Minute)

}

func (t *LocalTime) MarshalJSON() ([]byte, error) {
	tTime := time.Time(*t)
	if tTime.IsZero() {
		return []byte("\"\""), nil
	}
	return []byte(fmt.Sprintf("\"%v\"", tTime.Format("2006-01-02 15:04:05"))), nil
}

func (t LocalTime) Value() (driver.Value, error) {
	var zeroTime time.Time
	tlt := time.Time(t)
	//判断给定时间是否和默认零时间的时间戳相同
	if tlt.UnixNano() == zeroTime.UnixNano() {
		return nil, nil
	}
	return tlt, nil
}

func (t *LocalTime) Scan(v interface{}) error {
	if value, ok := v.(time.Time); ok {
		*t = LocalTime(value)
		return nil
	}
	return fmt.Errorf("can not convert %v to timestamp", v)
}