database/client.go

// Package mysql ...
//
// Description : mysql客户端
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2021-03-01 9:20 下午
package mysql

import (
	"errors"
	"fmt"
	"path/filepath"
	"strings"
	"sync"

	"git.zhangdeman.cn/zhangdeman/util"

	"git.zhangdeman.cn/zhangdeman/logger/wrapper"
	"github.com/gin-gonic/gin"

	"go.uber.org/zap"

	gormLogger "gorm.io/gorm/logger"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

var (
	// Client mysql客户端
	Client *client
)

func init() {
	Client = &client{
		lock:        &sync.RWMutex{},
		clientTable: make(map[string]*DBClient),
	}
}

type client struct {
	lock        *sync.RWMutex
	clientTable map[string]*DBClient
	logger      *zap.Logger
}

// AddWithConfigFile 使用文件生成新的客户端，文件名去掉后缀作为flag
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 19:19 2022/6/5
func (c *client) AddWithConfigFile(cfgFilePath string, logInstance *zap.Logger) error {
	var (
		err      error
		cfg      *cfgFile
		dbClient *DBClient
	)

	if cfg, err = c.getCfg(cfgFilePath); nil != err {
		return err
	}
	if nil == cfg {
		// 不支持的配置文件格式
		return nil
	}
	dbClient = &DBClient{
		dbFlag:         cfg.Flag,
		loggerInstance: logInstance,
		master:         nil,
		slave:          nil,
		extraFieldList: nil,
		cfg:            Mysql{},
	}
	if dbClient.master, err = c.GetDatabaseClient(cfg.Config.Master, logInstance); nil != err {
		return err
	}
	if dbClient.slave, err = c.GetDatabaseClient(cfg.Config.Slave, logInstance); nil != err {
		return err
	}
	c.lock.Lock()
	c.clientTable[dbClient.dbFlag] = dbClient
	c.lock.Unlock()
	return nil
}

// BatchAddWithConfigDir 自动读取目录下配置文件, 生成客户端
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 19:19 2022/6/5
func (c *client) BatchAddWithConfigDir(cfgDir string, logInstance *zap.Logger) error {
	filepathNames, _ := filepath.Glob(filepath.Join(cfgDir, "*"))
	for i := range filepathNames {
		if err := c.AddWithConfigFile(filepathNames[i], logInstance); nil != err {
			return err
		}
	}
	return nil
}

// getCfg 读取配置
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 18:05 2022/6/11
func (c *client) getCfg(cfgPath string) (*cfgFile, error) {
	fileArr := strings.Split(cfgPath, ".")
	if len(fileArr) < 2 {
		// 获取不到类型
		return nil, errors.New("文件格式必须是JSON或者YAML")
	}
	fileType := strings.ToLower(fileArr[len(fileArr)-1])
	fileFlagArr := strings.Split(fileArr[0], string(filepath.Separator))
	result := &cfgFile{
		Path:   cfgPath,
		Type:   "",
		Flag:   fileFlagArr[len(fileFlagArr)-1],
		Config: Database{},
	}
	var (
		err     error
		cfgInfo Database
	)
	switch fileType {
	case FileTypeYaml:
		fallthrough
	case FileTypeYml:
		result.Type = FileTypeYaml
		if err = util.File.ReadYmlContent(cfgPath, &result.Config); nil != err {
			return nil, fmt.Errorf("%s 配置文件解析失败, 原因 : %s", cfgPath, err.Error())
		}
	case FileTypeJson:
		result.Type = FileTypeJson
		if err = util.File.ReadJSONContent(cfgPath, &cfgInfo); nil != err {
			return nil, fmt.Errorf("%s 配置文件解析失败, 原因 : %s", cfgPath, err.Error())
		}
	default:
		// 不是JSON , 也不是YML, 跳过
		return nil, nil
	}
	if len(result.Config.Master.Timezone) == 0 {
		// 默认使用本地时区
		result.Config.Master.Timezone = "Local"
	} else {
		result.Config.Slave.Timezone = result.Config.Master.Timezone
	}
	return result, nil
}

// GetDBClient 获取db client
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 19:32 2022/6/5
func (c *client) GetDBClient(dbFlag string) (*DBClient, error) {
	c.lock.RLock()
	defer c.lock.RUnlock()
	var (
		exist    bool
		dbClient *DBClient
	)
	if dbClient, exist = c.clientTable[dbFlag]; !exist {
		return nil, fmt.Errorf("%s 标识的数据库实例不存在! ", dbFlag)
	}
	return dbClient, nil
}

// GetMasterClient 获取主库客户端
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 19:36 2022/6/5
func (c *client) GetMasterClient(ctx *gin.Context, dbFlag string) (*gorm.DB, error) {
	var (
		err      error
		dbClient *DBClient
	)
	if dbClient, err = c.GetDBClient(dbFlag); nil != err {
		return nil, err
	}

	session := dbClient.master.Session(&gorm.Session{})
	session.Logger = dbClient.getLogger(ctx, session, dbFlag+"-master")
	return session, nil
}

// GetSlaveClient 获取从库客户端
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 19:37 2022/6/5
func (c *client) GetSlaveClient(ctx *gin.Context, dbFlag string) (*gorm.DB, error) {
	var (
		err      error
		dbClient *DBClient
	)
	if dbClient, err = c.GetDBClient(dbFlag); nil != err {
		return nil, err
	}

	session := dbClient.slave.Session(&gorm.Session{})
	session.Logger = dbClient.getLogger(ctx, session, dbFlag+"-slave")
	return session, nil
}

// getGormClient 获取GORM client方法
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 11:24 2022/6/6
func (c *client) getGormClient() (*gorm.DB, error) {
	return nil, nil
}

// GetDatabaseClient 获取数据库连接
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 18:41 2022/6/11
func (c *client) GetDatabaseClient(conf *Mysql, logInstance *zap.Logger) (*gorm.DB, error) {
	var (
		instance *gorm.DB
		err      error
	)

	if instance, err = gorm.Open(mysql.Open(c.buildConnectionDSN(conf)), &gorm.Config{}); nil != err {
		return nil, err
	}

	instance.Logger = wrapper.NewGormLoggerWithInstance(nil, instance, logInstance, "", nil)

	return instance, nil
}

// buildConnectionDSN 构建连接信息
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 18:42 2022/6/11
func (c *client) buildConnectionDSN(conf *Mysql) string {
	return fmt.Sprintf(
		"%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=%s",
		conf.Username,
		conf.Password,
		conf.Host,
		conf.Port,
		conf.Database,
		conf.Charset,
		conf.Timezone,
	)
}

// DBClient 包装日志实例
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 3:09 PM 2021/12/24
type DBClient struct {
	dbFlag         string      // 数据库标识
	loggerInstance *zap.Logger // 日志实例
	master         *gorm.DB    // 主库
	slave          *gorm.DB    // 从库
	extraFieldList []string    // 提取的字段
	cfg            Mysql       // 数据库配置
}

// SetFlag 设置数据库标识
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 16:18 2022/6/5
func (dc *DBClient) SetFlag(dbFlag string) {
	dc.dbFlag = dbFlag
}

// GetMaster 获取主库连接
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 3:28 PM 2021/12/24
func (dc *DBClient) GetMaster(ctx *gin.Context) *gorm.DB {
	session := dc.master.Session(&gorm.Session{})
	session.Logger = dc.getLogger(ctx, session, "slave")
	return session
}

// GetSlave 获取从库链接
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 3:29 PM 2021/12/24
func (dc *DBClient) GetSlave(ctx *gin.Context) *gorm.DB {
	session := dc.slave.Session(&gorm.Session{})
	session.Logger = dc.getLogger(ctx, session, "slave")
	return session
}

// getLogger 获取日志实例
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 3:45 PM 2021/12/24
func (dc *DBClient) getLogger(ctx *gin.Context, dbClient *gorm.DB, node string) gormLogger.Interface {
	return wrapper.NewGormLoggerWithInstance(ctx, dbClient, dc.loggerInstance, dc.dbFlag+"|"+node, dc.extraFieldList)
}