server/components/gorm/component.go

package mhayaGORM

import (
	"fmt"
	"time"

	cfacade "github.com/mhaya/facade"
	clog "github.com/mhaya/logger"
	cprofile "github.com/mhaya/profile"
	"gorm.io/driver/clickhouse"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

const (
	Name             = "gorm_component"
	MysqlDBType      = "mysql"
	ClickhouseDBType = "clickhouse"
	mysqlDsn         = "%s:%s@(%s)/%s?charset=utf8&parseTime=True&loc=Local"
	clickhouseDsn    = "clickhouse://%s:%s@%s/%s?dial_timeout=10s&read_timeout=20s"
)

type (
	Component struct {
		cfacade.Component
		// key:groupId,value:{key:id,value:*gorm.Db}
		ormMap map[string]map[string]*gorm.DB
	}

	dbConfig struct {
		Enable         bool
		GroupId        string
		Id             string
		DbName         string
		Host           string
		UserName       string
		Password       string
		MaxIdleConnect int
		MaxOpenConnect int
		LogMode        bool
		DSN            string
		DBType         string
	}

	// HashDb hash by group id
	HashDb func(dbMaps map[string]*gorm.DB) string
)

func NewComponent() *Component {
	return &Component{
		ormMap: make(map[string]map[string]*gorm.DB),
	}
}

func (s *Component) Name() string {
	return Name
}

func parseDbConfig(groupId string, item cfacade.ProfileJSON) *dbConfig {
	return &dbConfig{
		GroupId:        groupId,
		Id:             item.GetString("db_id"),
		DBType:         item.GetString("db_type", ""),
		DSN:            item.GetString("dsn", ""),
		DbName:         item.GetString("db_name"),
		Host:           item.GetString("host"),
		UserName:       item.GetString("user_name"),
		Password:       item.GetString("password"),
		MaxIdleConnect: item.GetInt("max_idle_connect", 4),
		MaxOpenConnect: item.GetInt("max_open_connect", 8),
		LogMode:        item.GetBool("log_mode", true),
		Enable:         item.GetBool("enable", true),
	}
}

func (s *Component) Init() {
	// load only the database contained in the `db_id_list`
	dbIdList := s.App().Settings().Get("db_id_list")
	if dbIdList.LastError() != nil || dbIdList.Size() < 1 {
		clog.Warnf("[nodeId = %s] `db_id_list` property not exists.", s.App().NodeId())
		return
	}

	dbConfig := cprofile.GetConfig("db")
	if dbConfig.LastError() != nil {
		clog.Panic("`db` property not exists in profile file.")
	}

	for _, groupId := range dbConfig.Keys() {
		s.ormMap[groupId] = make(map[string]*gorm.DB)

		dbGroup := dbConfig.GetConfig(groupId)
		for i := 0; i < dbGroup.Size(); i++ {
			item := dbGroup.GetConfig(i)
			config := parseDbConfig(groupId, item)

			for _, key := range dbIdList.Keys() {
				if dbIdList.Get(key).ToString() != config.Id {
					continue
				}

				if !config.Enable {
					clog.Panicf("[dbName = %s] is disabled!", config.DbName)
				}

				switch config.DBType {
				case MysqlDBType:
					db, err := s.createORM(config)
					if err != nil {
						clog.Panicf("[dbName = %s] create orm fail. error = %s", config.DbName, err)
					}

					s.ormMap[groupId][config.Id] = db

				case ClickhouseDBType:
					db, err := s.createClickHouseORM(config)
					if err != nil {
						clog.Panicf("[dbName = %s] create chickorm fail. error = %s", config.DbName, err)
					}

					s.ormMap[groupId][config.Id] = db

				default:
					clog.Panicf("unknow dbName = %s", config.DbName)
				}

				clog.Infof("[dbGroup =%s, dbName = %s] is connected.", config.GroupId, config.Id)
			}
		}
	}
}

func (s *Component) createORM(cfg *dbConfig) (*gorm.DB, error) {
	db, err := gorm.Open(mysql.Open(cfg.GetDSN()), &gorm.Config{
		Logger: getLogger(),
	})

	if err != nil {
		return nil, err
	}

	sqlDB, err := db.DB()
	if err != nil {
		return nil, err
	}

	sqlDB.SetMaxIdleConns(cfg.MaxIdleConnect)
	sqlDB.SetMaxOpenConns(cfg.MaxOpenConnect)
	sqlDB.SetConnMaxLifetime(time.Minute)

	if cfg.LogMode {
		return db.Debug(), nil
	}

	return db, nil
}

func getLogger() logger.Interface {
	return logger.New(
		gormLogger{log: clog.DefaultLogger},
		logger.Config{
			SlowThreshold: time.Second,
			LogLevel:      logger.Silent,
			Colorful:      true,
		},
	)
}

func (s *Component) GetDb(id string) *gorm.DB {
	for _, group := range s.ormMap {
		for k, v := range group {
			if k == id {
				return v
			}
		}
	}
	return nil
}

func (s *Component) GetHashDb(groupId string, hashFn HashDb) (*gorm.DB, bool) {
	dbGroup, found := s.GetDbMap(groupId)
	if !found {
		clog.Warnf("groupId = %s not found.", groupId)
		return nil, false
	}

	dbId := hashFn(dbGroup)
	db, found := dbGroup[dbId]
	return db, found
}

func (s *Component) GetDbMap(groupId string) (map[string]*gorm.DB, bool) {
	dbGroup, found := s.ormMap[groupId]
	return dbGroup, found
}

func (s *dbConfig) GetDSN() string {
	if s.DSN != "" {
		return s.DSN
	}

	switch s.DBType {
	case MysqlDBType:
		s.DSN = mysqlDsn
		return fmt.Sprintf(s.DSN, s.UserName, s.Password, s.Host, s.DbName)

	case ClickhouseDBType:
		s.DSN = clickhouseDsn
		return fmt.Sprintf(s.DSN, s.UserName, s.Password, s.Host, s.DbName)

	default:
		// never reach
		return ""
	}
}

func (s *Component) createClickHouseORM(cfg *dbConfig) (*gorm.DB, error) {
	db, err := gorm.Open(clickhouse.Open(cfg.GetDSN()), &gorm.Config{
		Logger: getLogger(),
	})
	if err != nil {
		return nil, err
	}

	sqlDB, err := db.DB()
	if err != nil {
		return nil, err
	}
	sqlDB.SetConnMaxLifetime(time.Minute)

	if cfg.LogMode {
		return db.Debug(), nil
	}

	return db, nil
}