gopkgsv2/database/init.go

// Copyright 2019 getensh.com. All rights reserved.
// Use of this source code is governed by getensh.com.

package database

import (
	"fmt"
	"git.getensh.com/common/gopkgsv2/logger"
	"go.uber.org/zap"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"log"
	"os"
	"time"
)

var (
	db *gorm.DB
)

// Setup 建立连接
func Setup(user, passwd, addr, dbname, charset string, maxIdle, maxConn int, logMode bool) *gorm.DB {
	conf := Config{
		SlowThreshold:             time.Second,
		Colorful:                  true,
		IgnoreRecordNotFoundError: false,
		LogLevel:                  zap.ErrorLevel,
		TraceWithLevel:            zap.ErrorLevel,
		Zap:                       logger.Logger,
	}
	// 是否开启调试模式
	if logMode {
		conf.LogLevel = zap.DebugLevel
		conf.TraceWithLevel = zap.DebugLevel
		conf.SlowThreshold = time.Millisecond * 200
	}
	config := &gorm.Config{
		Logger: NewGormLogger(log.New(os.Stdout, "\r\n", log.LstdFlags), conf),
	}
	// 组装参数
	dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=%s&parseTime=True&loc=Local",
		user, passwd, addr, dbname, charset)
	// 打开新连接
	var err error
	db, err = gorm.Open(mysql.Open(dsn), config)
	if err != nil {
		log.Fatal("open mysql connection failed. err: ", err)
	}

	// 其他设置

	sqlDB, err := db.DB()
	if err != nil {
		log.Fatal("open mysql connection failed. err: ", err)
	}
	sqlDB.SetMaxIdleConns(maxIdle)
	sqlDB.SetMaxOpenConns(maxConn)
	RegisterCallback(db)
	return db
}

// DB 获取连接
func DB() *gorm.DB {
	return db
}

// Close 关闭连接
func Close() {
	if db != nil {
		sqlDB, _ := db.DB()
		_ = sqlDB.Close()
	}
}