package gfdb import ( "context" "database/sql" "encoding/json" "fmt" "regexp" "strings" "time" "gitea.com/red-future/common/utils" "github.com/bwmarrin/snowflake" "github.com/gogf/gf/v2/crypto/gmd5" "github.com/gogf/gf/v2/database/gdb" "github.com/gogf/gf/v2/database/gredis" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/os/gcache" "github.com/gogf/gf/v2/os/glog" "github.com/gogf/gf/v2/text/gstr" "github.com/gogf/gf/v2/util/gconv" "go.opentelemetry.io/otel/trace" ) // ==================== 缓存管理器(单例) ==================== var ( localCache *gcache.Cache ) // getLocalCache 获取本地缓存实例 func getLocalCache() *gcache.Cache { if localCache == nil { localCache = gcache.New() } return localCache } // getFromCache 从缓存获取数据(本地缓存 -> Redis) func getFromCache(ctx context.Context, key string) ([]byte, bool) { // 1. 先查本地缓存 if val, err := getLocalCache().Get(ctx, key); err == nil && val != nil { if data := val.Bytes(); len(data) > 0 { return data, true } } // 2. 再查Redis缓存 if g.Redis() != nil { result, err := g.Redis().Get(ctx, key) if err == nil && !result.IsEmpty() { data := result.Bytes() // 写入本地缓存 err = getLocalCache().Set(ctx, key, data, time.Duration(g.Cfg().MustGet(ctx, "cache.localTTL").Int64())*time.Second) if err != nil { return nil, false } return data, true } } return nil, false } // setToCache 写入缓存(本地缓存 + Redis) func setToCache(ctx context.Context, key string, data []byte) (err error) { if len(data) == 0 { return } // 1. 写入本地缓存 if err = getLocalCache().Set(ctx, key, data, time.Duration(g.Cfg().MustGet(ctx, "cache.localTTL").Int64())*time.Second); err != nil { return } // 2. 写入Redis缓存 if g.Redis() != nil { _, err = g.Redis().Set(ctx, key, data, gredis.SetOption{ TTLOption: gredis.TTLOption{ EX: gconv.PtrInt64(g.Cfg().MustGet(ctx, "cache.redisTTL")), }, }) if err != nil { return } } return } // deleteCacheByPattern 根据模式删除缓存 func deleteCacheByPattern(ctx context.Context, pattern string) (err error) { // 1. 删除匹配模式的本地缓存 localCache := getLocalCache() keys := localCache.MustKeyStrings(ctx) if len(keys) > 0 { for _, key := range keys { if matchPattern(key, pattern) { _, err = localCache.Remove(ctx, key) if err != nil { return err } } } } // 2. 删除Redis缓存(使用SCAN+DEL) if g.Redis() != nil { keys, err := g.Redis().Keys(ctx, pattern) if err != nil { return err } for _, key := range keys { _, err = g.Redis().Del(ctx, key) if err != nil { return err } } } return nil } // matchPattern 检查 key 是否匹配 Redis SCAN 的 MATCH 模式(支持 * 通配符) func matchPattern(key string, pattern string) bool { // 将 Redis 的 MATCH 模式转换为正则表达式 // 转义正则特殊字符(除了 *) regexPattern := regexp.QuoteMeta(pattern) // 将转义后的 \* 替换回 .* regexPattern = strings.ReplaceAll(regexPattern, `\*`, ".*") // 添加开始和结束锚点 regexPattern = "^" + regexPattern + "$" matched, _ := regexp.MatchString(regexPattern, key) return matched } // ==================== 统一Hook入口 ==================== // CatchSQLHook 返回统一的 HookHandler(包含租户自动赋值和缓存) // 使用示例: // // // 基础使用(自动租户赋值,无缓存) // g.DB().Model("user").Hook(base.CatchSQLHook()).Ctx(ctx).Insert(data) // // // 启用缓存(用户无感知,自动处理缓存key) // ctx = base.WithCacheEnabled(ctx, "asset") // Asset.CtxWithCache(ctx).Where("id", 123).Scan(&result) func catchSQLHook() gdb.HookHandler { return gdb.HookHandler{ Insert: insertHook, Update: updateHook, Delete: deleteHook, Select: selectHook, } } // ==================== Insert钩子 ==================== func insertHook(ctx context.Context, in *gdb.HookInsertInput) (result sql.Result, err error) { userInfo, err := utils.GetUserInfo(ctx) if err != nil { return nil, err } node, err := snowflake.NewNode(g.Cfg().MustGet(ctx, "server.workerId").Int64()) if err != nil { return nil, err } for i := range in.Data { if _, ok := in.Data[i]["id"]; ok { in.Data[i]["id"] = node.Generate().Int64() } if _, ok := in.Data[i]["tenant_id"]; ok { if !g.IsEmpty(userInfo.TenantId) { in.Data[i]["tenant_id"] = userInfo.TenantId } else { return nil, fmt.Errorf("tenantId cannot be empty") } } if _, ok := in.Data[i]["creator"]; ok { if !g.IsEmpty(userInfo.UserName) { in.Data[i]["creator"] = userInfo.UserName } else { return nil, fmt.Errorf("user info cannot be empty") } } if _, ok := in.Data[i]["updater"]; ok { if !g.IsEmpty(userInfo.UserName) { in.Data[i]["updater"] = userInfo.UserName } else { return nil, fmt.Errorf("user info cannot be empty") } } } // 2. 执行插入 result, err = in.Next(ctx) if err != nil { return nil, err } // 3. 清除相关缓存 if userInfo != nil && userInfo.TenantId != 0 { if err = deleteCacheByPattern(ctx, getCacheKey(userInfo.TenantId, in.Table, true)); err != nil { return nil, err } } return result, nil } // ==================== Update钩子 ==================== func updateHook(ctx context.Context, in *gdb.HookUpdateInput) (result sql.Result, err error) { // 1. 自动赋值修改人 userInfo, _ := utils.GetUserInfo(ctx) switch data := in.Data.(type) { case gdb.Map: if _, ok := data["updater"]; ok { if !g.IsEmpty(userInfo.UserName) { data["updater"] = userInfo.UserName } else { return nil, fmt.Errorf("user info cannot be empty") } } case gdb.List: for i := range data { if !g.IsEmpty(userInfo.UserName) { if _, ok := data[i]["updater"]; ok { if !g.IsEmpty(userInfo.UserName) { data[i]["updater"] = userInfo.UserName } else { return nil, fmt.Errorf("user info cannot be empty") } } } } } // 2. 执行更新 result, err = in.Next(ctx) if err != nil { return nil, err } // 3. 清除相关缓存 if userInfo != nil && userInfo.TenantId != 0 { if err = deleteCacheByPattern(ctx, getCacheKey(userInfo.TenantId, in.Table, true)); err != nil { return nil, err } } return result, nil } // ==================== Delete钩子 ==================== func deleteHook(ctx context.Context, in *gdb.HookDeleteInput) (result sql.Result, err error) { // 1. 执行删除 result, err = in.Next(ctx) if err != nil { return nil, err } // 2. 清除相关缓存 userInfo, _ := utils.GetUserInfo(ctx) if userInfo != nil && userInfo.TenantId != 0 { if err = deleteCacheByPattern(ctx, getCacheKey(userInfo.TenantId, in.Table, true)); err != nil { return nil, err } } return result, nil } // ==================== Select钩子(缓存读取) ==================== func selectHook(ctx context.Context, in *gdb.HookSelectInput) (result gdb.Result, err error) { traceID := getTraceID(ctx) enabled, err := gcache.Get(ctx, traceID) // 未启用缓存,直接执行查询 if !gconv.Bool(enabled) { return in.Next(ctx) } // 从 SQL 字符串中提取 WHERE 条件部分 whereCondition := "" // 查找 WHERE 关键字(不区分大小写) whereIndex := gstr.PosI(in.Sql, " WHERE ") if whereIndex != -1 { // 提取 WHERE 之后的内容 whereCondition = in.Sql[whereIndex+7:] // 移除 ORDER BY, GROUP BY, HAVING, LIMIT 等后续子句 for _, keyword := range []string{" ORDER BY ", " GROUP BY ", " HAVING ", " LIMIT ", " FOR UPDATE"} { if idx := gstr.PosI(whereCondition, keyword); idx != -1 { whereCondition = whereCondition[:idx] } } } user, err := utils.GetUserInfo(ctx) if err != nil { return nil, err } encrypt, err := gmd5.Encrypt(fmt.Sprintf("%s:%s", whereCondition, in.Args)) if err != nil { return nil, err } // 构建缓存key:sql:tenantId:table:where条件:args cacheKey := fmt.Sprintf("%s:%s:%s", getCacheKey(user.TenantId, in.Table, false), getSelectTypeString(in.SelectType), encrypt) // 1. 先查缓存 if data, ok := getFromCache(ctx, cacheKey); ok { var records gdb.Result if err := json.Unmarshal(data, &records); err == nil && len(records) > 0 { return records, nil } } // 2. 执行数据库查询 result, err = in.Next(ctx) if err != nil { return nil, err } // 3. 写入缓存 if len(result) > 0 { if data, err := json.Marshal(result); err == nil { if err = setToCache(ctx, cacheKey, data); err != nil { return nil, err } } } return result, nil } func getCacheKey(tenantId uint64, table string, isBlur bool) string { cacheKey := fmt.Sprintf("sql:tenantId-%v:%s", tenantId, table) if isBlur { cacheKey = fmt.Sprintf("%s:*", cacheKey) } return cacheKey } // getSelectTypeString 将 SelectType 枚举转换为可读字符串 func getSelectTypeString(selectType gdb.SelectType) string { switch selectType { case gdb.SelectTypeDefault: return "default" case gdb.SelectTypeCount: return "count" case gdb.SelectTypeValue: return "value" case gdb.SelectTypeArray: return "array" default: return "unknown" } } // ==================== 调用方法 ==================== var ( schemaPrefix = "tenant-" ) type Gfdb interface { Exec(ctx context.Context, sql string, args ...any) (sql.Result, error) Model(ctx context.Context, tableNameOrStruct ...any) *model Transaction(ctx context.Context, f func(ctx context.Context, tx gdb.TX) error) error } type cache interface { Cache(ctx context.Context) *gdb.Model } type model struct { *gdb.Model } type dataBase struct { gdb.DB } func GetTablePrefix(ctx context.Context) (prefix string, err error) { tenantId, config, err := checkSchemaConfig(ctx) if err != nil { glog.Errorf(ctx, "[DB] checkSchemaConfig error: %v", err) return } if config { sprintf := fmt.Sprintf("database.%s%v.0.prefix", schemaPrefix, tenantId) prefix = g.Cfg().MustGet(ctx, sprintf).String() return } prefix = g.Cfg().MustGet(ctx, "database.default.0.prefix").String() return } func checkSchemaConfig(ctx context.Context) (uint64, bool, error) { user, err := utils.GetUserInfo(ctx) if err != nil { glog.Errorf(ctx, "[DB] GetUserInfo error: %v", err) return 0, false, err } var schema = fmt.Sprintf("%s%v", schemaPrefix, user.TenantId) sprintf := fmt.Sprintf("database.%s", schema) if !g.Cfg().MustGet(ctx, sprintf).IsEmpty() { return user.TenantId, true, nil } return user.TenantId, false, nil } func DB(ctx context.Context) Gfdb { tenantId, config, err := checkSchemaConfig(ctx) if err != nil { glog.Errorf(ctx, "[DB] checkSchemaConfig error: %v", err) return nil } var schema = fmt.Sprintf("%s%v", schemaPrefix, tenantId) var dbName []string if config { dbName = append(dbName, schema) } else { dbName = append(dbName, "default") // 配置文件中 default 是数组格式,需要通过索引 0 访问 defaultConfig := g.Cfg().MustGet(ctx, "database.default") if defaultConfig.IsSlice() { schema = g.Cfg().MustGet(ctx, "database.default.0.name").String() } else { schema = g.Cfg().MustGet(ctx, "database.default.name").String() } } return &dataBase{ DB: g.DB(dbName...).Schema(schema), } } func (d *dataBase) Model(ctx context.Context, tableNameOrStruct ...any) *model { m := d.DB.Model(tableNameOrStruct...).Ctx(ctx) tenantId, config, err := checkSchemaConfig(ctx) if err != nil { glog.Errorf(ctx, "[DB] checkSchemaConfig error: %v", err) return nil } if config { // 创建按地区分库的配置 shardingConfig := gdb.ShardingConfig{ Schema: gdb.ShardingSchemaConfig{ Enable: true, // 启用分库 Prefix: schemaPrefix, // 分库前缀 Rule: &RegionShardingRule{RegionMapping: tenantId}, // 自定义分库规则 }, } m.Sharding(shardingConfig).ShardingValue(tenantId) } m.OmitNil().Hook(catchSQLHook()) return &model{ Model: m, } } func (d *dataBase) Transaction(ctx context.Context, f func(ctx context.Context, tx gdb.TX) error) error { return d.DB.Transaction(ctx, f) } func (d *model) Cache(ctx context.Context) *gdb.Model { traceID := getTraceID(ctx) if traceID == "" { glog.Errorf(ctx, "[DB] GetTraceID error: traceID is empty") return nil } if err := gcache.Set(ctx, traceID, true, time.Second); err != nil { glog.Errorf(ctx, "[DB] Cache error: %v", err) return nil } return d.Model } // getTraceID 从 context 中获取链路追踪 ID func getTraceID(ctx context.Context) string { span := trace.SpanFromContext(ctx) if span != nil && span.SpanContext().HasTraceID() { return span.SpanContext().TraceID().String() } return "" } type RegionShardingRule struct { RegionMapping uint64 } func (r *RegionShardingRule) SchemaName(ctx context.Context, config gdb.ShardingSchemaConfig, value any) (string, error) { region, ok := value.(uint64) if !ok { return "", fmt.Errorf("sharding value must be string for RegionShardingRule") } if r.RegionMapping == region { return config.Prefix + gconv.String(region), nil } return "default", nil } // TableName 实现分表规则接口 func (r *RegionShardingRule) TableName(ctx context.Context, config gdb.ShardingTableConfig, value any) (string, error) { // 这里不实现分表,返回空字符串 return "", nil }