OAuth-core/tools/database/config.go

package database

import (
	"time"

	"gorm.io/gorm"
	"gorm.io/plugin/dbresolver"
)

type DBConfig struct {
	dsn             string
	connMaxIdleTime int
	connMaxLifetime int
	maxIdleConns    int
	maxOpenConns    int
	registers       []ResolverConfigure
}

// NewConfigure 初始化 Configure
func NewConfigure(
	dsn string,
	maxIdleConns,
	maxOpenConns,
	connMaxIdleTime,
	connMaxLifetime int,
	registers []ResolverConfigure) Configure {
	return &DBConfig{
		dsn:             dsn,
		connMaxIdleTime: connMaxIdleTime,
		connMaxLifetime: connMaxLifetime,
		maxIdleConns:    maxIdleConns,
		maxOpenConns:    maxOpenConns,
		registers:       registers,
	}
}

// Init 获取db，⚠️注意：读写分离只能配置一组
func (e *DBConfig) Init(config *gorm.Config, open func(string) gorm.Dialector) (*gorm.DB, error) {
	db, err := gorm.Open(open(e.dsn), config)
	if err != nil {
		return nil, err
	}
	var register *dbresolver.DBResolver
	for i := range e.registers {
		register = e.registers[i].Init(register, open)
	}
	if register == nil {
		register = dbresolver.Register(dbresolver.Config{})
	}
	if e.connMaxIdleTime > 0 {
		register = register.SetConnMaxIdleTime(time.Duration(e.connMaxIdleTime) * time.Second)
	}
	if e.connMaxLifetime > 0 {
		register = register.SetConnMaxLifetime(time.Duration(e.connMaxLifetime) * time.Second)
	}
	if e.maxOpenConns > 0 {
		register = register.SetMaxOpenConns(e.maxOpenConns)
	}
	if e.maxIdleConns > 0 {
		register = register.SetMaxIdleConns(e.maxIdleConns)
	}
	if register != nil {
		err = db.Use(register)
	}
	return db, err
}

type DBResolverConfig struct {
	sources  []string
	replicas []string
	policy   string
	tables   []interface{}
}

// NewResolverConfigure 初始化 ResolverConfigure
func NewResolverConfigure(sources, replicas []string, policy string, tables []string) ResolverConfigure {
	data := make([]interface{}, len(tables))
	for i := range tables {
		data[i] = tables[i]
	}
	return &DBResolverConfig{
		sources:  sources,
		replicas: replicas,
		policy:   policy,
		tables:   data,
	}
}

func (e *DBResolverConfig) Init(
	register *dbresolver.DBResolver,
	open func(string) gorm.Dialector) *dbresolver.DBResolver {
	if len(e.tables) == 0 && len(e.sources) == 0 && len(e.replicas) == 0 {
		return register
	}
	var config dbresolver.Config
	if len(e.sources) > 0 {
		config.Sources = make([]gorm.Dialector, len(e.sources))
		for i := range e.sources {
			config.Sources[i] = open(e.sources[i])
		}
	}
	if len(e.replicas) > 0 {
		config.Replicas = make([]gorm.Dialector, len(e.replicas))
		for i := range e.replicas {
			config.Replicas[i] = open(e.replicas[i])
		}
	}
	if e.policy != "" {
		policy, ok := policies[e.policy]
		if ok {
			config.Policy = policy
		}
	}
	if register == nil {
		register = dbresolver.Register(config, e.tables...)
		return register
	}
	register = register.Register(config, e.tables...)
	return register
}