// 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() }