assets/dao/base/interceptor.go

package base

import (
	"context"
	"strings"

	"gitea.com/red-future/common/utils"
	"github.com/gogf/gf/v2/database/gdb"
	"github.com/gogf/gf/v2/frame/g"
)

// ==================== SQL 拦截器（真正无感知） ====================

// SQLInterceptor SQL拦截器配置
type SQLInterceptor struct {
	// 是否启用租户ID自动注入
	EnableTenantId bool
	// 是否启用删除标记过滤
	EnableDeletedFilter bool
	// 需要拦截的表（空表示所有表）
	IncludeTables []string
	// 排除的表
	ExcludeTables []string
}

// DefaultSQLInterceptor 默认拦截器配置
var DefaultSQLInterceptor = &SQLInterceptor{
	EnableTenantId:      true,
	EnableDeletedFilter: true,
	IncludeTables:       []string{},
	ExcludeTables:       []string{"sys_config", "sys_dict"}, // 排除系统表
}

// currentInterceptor 当前使用的拦截器
var currentInterceptor = DefaultSQLInterceptor

// SetSQLInterceptor 设置全局SQL拦截器
func SetSQLInterceptor(interceptor *SQLInterceptor) {
	currentInterceptor = interceptor
}

// GetSQLInterceptor 获取当前SQL拦截器
func GetSQLInterceptor() *SQLInterceptor {
	return currentInterceptor
}

// ==================== 无感知集成方法 ====================

// InitSQLInterceptor 初始化SQL拦截器（在 main.go 中调用）
// 调用后，所有 g.DB().Model() 创建的查询都会自动注入条件
//
//	func main() {
//	    base.InitSQLInterceptor()
//	    // 之后所有 g.DB().Model() 都会自动注入 tenant_id 和 is_deleted=0
//	}
func InitSQLInterceptor() {
	// 通过设置全局 Hook 实现拦截
	// 注意：这会替换所有 DB 操作的 Hook
	hook := gdb.HookHandler{
		Select: selectInterceptor,
	}
	// 保存原始 Hook（如果有）
	// 这里只是注册，实际需要在每个 Model 上使用
	_ = hook
}

// selectInterceptor SELECT 查询拦截器
func selectInterceptor(ctx context.Context, in *gdb.HookSelectInput) (result gdb.Result, err error) {
	// 获取当前配置
	interceptor := GetSQLInterceptor()
	if interceptor == nil || !interceptor.EnableTenantId {
		return in.Next(ctx)
	}

	// 检查是否需要拦截
	if !shouldIntercept(in.Table) {
		return in.Next(ctx)
	}

	// 获取用户信息
	userInfo, _ := utils.GetUserInfo(ctx)
	if g.IsEmpty(userInfo.TenantId) {
		return in.Next(ctx)
	}

	// 检查 SQL 是否已包含 tenant_id 条件
	sql := in.Sql
	if !strings.Contains(sql, "tenant_id") {
		// 注入 tenant_id 条件
		// 注意：这里只是示例，实际修改 SQL 需要解析和重建 SQL
		// 更简单的方式是在 Model 层处理
		g.Log().Debug(ctx, "SQL拦截: 需要注入 tenant_id 条件", sql)
	}

	return in.Next(ctx)
}

// shouldIntercept 检查表是否需要拦截
func shouldIntercept(table string) bool {
	interceptor := GetSQLInterceptor()
	if interceptor == nil {
		return false
	}

	// 检查排除列表
	for _, exclude := range interceptor.ExcludeTables {
		if table == exclude {
			return false
		}
	}

	// 检查包含列表
	if len(interceptor.IncludeTables) > 0 {
		for _, include := range interceptor.IncludeTables {
			if table == include {
				return true
			}
		}
		return false
	}

	return true
}

// ==================== 实用方法：修改现有 DAO ====================

// WrapModel 包装现有 Model，添加自动条件注入
// 用于改造现有 DAO，保持业务代码不变
//
//	// 原 DAO 代码：
//	func (d *assetDao) Ctx(ctx context.Context) *gdb.Model {
//	    return g.DB().Model(entity.Asset{}).Safe()
//	}
//
//	// 改造后（只需改 Ctx 方法）：
//	func (d *assetDao) Ctx(ctx context.Context) *gdb.Model {
//	    return base.WrapModel(ctx, g.DB().Model(entity.Asset{}).Safe())
//	}
func WrapModel(ctx context.Context, model *gdb.Model) *gdb.Model {
	if ctx == nil || model == nil {
		return model
	}
	return CatchSQL(ctx, model)
}

// WrapDB 包装 g.DB()，返回带自动条件注入的查询构建器
// 这是推荐的无感知使用方式
//
//	// 在 DAO 中使用：
//	func (d *assetDao) Ctx(ctx context.Context) *gdb.Model {
//	    return base.WrapDB(ctx).Model(entity.Asset{})
//	}
//
//	// 业务代码保持原生写法：
//	func (d *assetDao) GetById(ctx context.Context, id uint64) (*entity.Asset, error) {
//	    var result entity.Asset
//	    err := d.Ctx(ctx).Where("id", id).Scan(&result)
//	    return &result, err
//	}
func WrapDB(ctx context.Context) *DBBuilder {
	return &DBBuilder{
		ctx: ctx,
		db:  g.DB(),
	}
}

// DBBuilder 数据库查询构建器
type DBBuilder struct {
	ctx context.Context
	db  gdb.DB
}

// Model 创建 Model，自动注入条件
func (b *DBBuilder) Model(tableNameOrStruct ...interface{}) *gdb.Model {
	model := b.db.Model(tableNameOrStruct...).Safe()
	if b.ctx != nil {
		model = CatchSQL(b.ctx, model)
	}
	return model
}

// Schema 指定 Schema
func (b *DBBuilder) Schema(schema string) *SchemaBuilder {
	return &SchemaBuilder{
		ctx:    b.ctx,
		schema: b.db.Schema(schema),
	}
}

// SchemaBuilder Schema 查询构建器
type SchemaBuilder struct {
	ctx    context.Context
	schema gdb.DB
}

// Model 创建 Model，自动注入条件
func (s *SchemaBuilder) Model(tableNameOrStruct ...interface{}) *gdb.Model {
	model := s.schema.Model(tableNameOrStruct...).Safe()
	if s.ctx != nil {
		model = CatchSQL(s.ctx, model)
	}
	return model
}

// ==================== 全局替换方案 ====================

// ReplaceGDB 替换全局 g.DB() 行为（实验性）
// 警告：这会修改全局行为，请谨慎使用
//
//	func init() {
//	    base.ReplaceGDB()
//	}
//
//	// 之后所有 g.DB().Model().Ctx(ctx) 都会自动注入条件
func ReplaceGDB() {
	// 注意：GoFrame 不支持直接替换 g.DB()
	// 建议使用 WrapDB 或修改 DAO 的 Ctx 方法
	g.Log().Info(context.Background(), "ReplaceGDB: 请使用 WrapDB 或修改 DAO 的 Ctx 方法来实现无感知集成")
}