im_manager/server/internal/library/imtenant/db.go

// Package imtenant provides helpers for accessing tenant business databases.
package imtenant

import (
	"context"
	"fmt"
	"hash/fnv"
	"hotgo/internal/dao"
	"hotgo/internal/model/entity"
	"net/url"
	"strings"
	"sync"

	_ "github.com/gogf/gf/contrib/drivers/mysql/v2"
	"github.com/gogf/gf/v2/database/gdb"
	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/frame/g"
)

var tenantDbGroups sync.Map

const registryDbGroup = "imRegistry"

var registryDbGroupOnce sync.Once

// DB resolves a tenant registry record and returns the corresponding business DB.
func DB(ctx context.Context, tenantId int64) (db gdb.DB, tenant *entity.ImTenant, err error) {
	if tenantId <= 0 {
		return nil, nil, gerror.New("租户ID不能为空")
	}

	if err = RegistryDB().Model(dao.ImTenant.Table()).Safe().Ctx(ctx).WherePri(tenantId).Scan(&tenant); err != nil {
		return nil, nil, gerror.Wrap(err, "获取租户信息失败")
	}
	if tenant == nil || tenant.Id <= 0 {
		return nil, nil, gerror.New("租户不存在")
	}

	link, err := tenantLink(ctx, tenant)
	if err != nil {
		return nil, nil, err
	}

	group := fmt.Sprintf("im_tenant_%d_%x", tenant.Id, hash(link))
	if cached, ok := tenantDbGroups.Load(group); ok && cached == link {
		db, err = gdb.Instance(group)
		return db, tenant, err
	}

	if err = gdb.SetConfigGroup(group, gdb.ConfigGroup{{Link: link}}); err != nil {
		return nil, nil, gerror.Wrap(err, "初始化租户业务库连接失败")
	}
	tenantDbGroups.Store(group, link)

	db, err = gdb.Instance(group)
	if err != nil {
		return nil, nil, gerror.Wrap(err, "获取租户业务库连接失败")
	}
	return db, tenant, nil
}

// RegistryDB returns the dedicated tenant registry connection, or the default
// HotGo database when a separate registry group is not configured.
func RegistryDB() gdb.DB {
	registryDbGroupOnce.Do(func() {
		ensureRegistryDBGroup()
	})
	if configs, err := gdb.GetConfigGroup(registryDbGroup); err == nil && len(configs) > 0 {
		return g.DB(registryDbGroup)
	}
	return g.DB()
}

func ensureRegistryDBGroup() {
	if configs, err := gdb.GetConfigGroup(registryDbGroup); err == nil && len(configs) > 0 {
		configs[0].Prefix = ""
		_ = gdb.SetConfigGroup(registryDbGroup, configs)
		return
	}

	configs, err := gdb.GetConfigGroup("default")
	if err != nil || len(configs) == 0 {
		return
	}
	cloned := make(gdb.ConfigGroup, len(configs))
	copy(cloned, configs)
	for i := range cloned {
		cloned[i].Prefix = ""
	}
	_ = gdb.SetConfigGroup(registryDbGroup, cloned)
}

func tenantLink(ctx context.Context, tenant *entity.ImTenant) (string, error) {
	if tenant.DbDsn != "" {
		return normalizeMySQLTimeOptions(tenant.DbDsn)
	}
	if tenant.DbName == "" {
		return "", gerror.New("租户业务库未配置")
	}

	defaultLink := defaultDbLink(ctx)
	if defaultLink == "" {
		return "", gerror.New("默认数据库连接未配置，无法推导租户业务库连接")
	}

	link, err := replaceMySQLDatabase(defaultLink, tenant.DbName)
	if err != nil {
		return "", err
	}
	return link, nil
}

func defaultDbLink(ctx context.Context) string {
	if configs, err := gdb.GetConfigGroup("default"); err == nil && len(configs) > 0 {
		if configs[0].Link != "" {
			return configs[0].Link
		}
	}

	cfg := g.Cfg()
	if !cfg.Available(ctx) {
		return ""
	}
	defaultLink := cfg.MustGet(ctx, "database.default.0.link").String()
	if defaultLink != "" {
		return defaultLink
	}
	return cfg.MustGet(ctx, "database.default.link").String()
}

func replaceMySQLDatabase(link string, dbName string) (string, error) {
	if !strings.HasPrefix(link, "mysql:") {
		return "", gerror.New("仅支持从 MySQL 默认连接推导租户业务库连接")
	}

	dbStart := strings.LastIndex(link, ")/")
	if dbStart < 0 {
		return "", gerror.New("默认数据库连接格式不正确")
	}
	dbStart += 2
	queryStart := strings.Index(link[dbStart:], "?")
	if queryStart < 0 {
		return normalizeMySQLTimeOptions(link[:dbStart] + dbName)
	}
	queryStart += dbStart
	return normalizeMySQLTimeOptions(link[:dbStart] + dbName + link[queryStart:])
}

func normalizeMySQLTimeOptions(link string) (string, error) {
	if !strings.HasPrefix(link, "mysql:") {
		if !strings.Contains(link, "@tcp(") {
			return "", gerror.New("租户业务库仅支持 MySQL")
		}
		link = "mysql:" + link
	}
	parts := strings.SplitN(link, "?", 2)
	values := make(url.Values)
	if len(parts) == 2 {
		var err error
		values, err = url.ParseQuery(parts[1])
		if err != nil {
			return "", gerror.Wrap(err, "租户业务库连接参数格式不正确")
		}
	}
	values.Set("loc", "UTC")
	values.Set("parseTime", "true")
	return parts[0] + "?" + values.Encode(), nil
}

func hash(s string) uint64 {
	h := fnv.New64a()
	_, _ = h.Write([]byte(s))
	return h.Sum64()
}