Blizzard
30d667954b
统一强制、别靠人肉：受租户模型标记 isTenantScoped() 后，store/tenant_scope.go 的 gorm 回调按请求 ctx 自动给查询加 WHERE tenant_id、创建自动填 tenant_id。 - KB / Agent 加 TenantID 字段 + isTenantScoped() 标记 - middleware.TenantContext 把 tenant 注入 request context 供 store 插件读取 - ctx 无租户（系统/回填/未登录）不过滤，保留跨租户操作能力 - 启动 BackfillRowTenants 回填存量行 tenant_id=owner 默认租户（幂等） live 验证：创建自动写 tenant_id ✓；同 owner 不同 tenant 的行被查询过滤 ✓。 Doc/DocLink（异步入库）、Task/Eval（无 owner）留待增量2b。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
92 lines
3.4 KiB
Go
92 lines
3.4 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"reflect"
|
|
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
)
|
|
|
|
// tenantScopedMarker：模型实现此(空)接口即"受租户隔离"——gorm tenant 插件据此对其
|
|
// 查询/更新/删除自动加 `tenant_id = <ctx租户>`、创建时自动填 tenant_id。
|
|
// 未实现的模型（User/Tenant/TenantMember/Prompt/Model/Pricing…）不受影响。
|
|
// 未实现方法用小写→只有本包模型能标记（KB/Agent 已标；Doc/DocLink/Task 待增量铺开）。
|
|
type tenantScopedMarker interface{ isTenantScoped() }
|
|
|
|
// ---- 请求上下文携带 tenant_id（中间件注入 → 传到 store → 插件读取）----
|
|
type ctxKeyTenant struct{}
|
|
|
|
// WithTenant 把 tenant_id 放进 context（空则原样返回，避免污染系统/回填查询）。
|
|
func WithTenant(ctx context.Context, tenantID string) context.Context {
|
|
if tenantID == "" {
|
|
return ctx
|
|
}
|
|
return context.WithValue(ctx, ctxKeyTenant{}, tenantID)
|
|
}
|
|
|
|
func tenantFromCtx(ctx context.Context) string {
|
|
if ctx == nil {
|
|
return ""
|
|
}
|
|
s, _ := ctx.Value(ctxKeyTenant{}).(string)
|
|
return s
|
|
}
|
|
|
|
// registerTenantScope 挂 gorm 回调，实现"统一强制、别靠人肉"的租户隔离。
|
|
func registerTenantScope(db *gorm.DB) {
|
|
q := func(d *gorm.DB) { addTenantWhere(d) }
|
|
_ = db.Callback().Query().Before("gorm:query").Register("tenant:query", q)
|
|
_ = db.Callback().Update().Before("gorm:update").Register("tenant:update", q)
|
|
_ = db.Callback().Delete().Before("gorm:delete").Register("tenant:delete", q)
|
|
_ = db.Callback().Create().Before("gorm:create").Register("tenant:create", setTenantOnCreate)
|
|
}
|
|
|
|
func isTenantScopedStmt(db *gorm.DB) bool {
|
|
if db.Statement.Schema == nil {
|
|
return false
|
|
}
|
|
_, ok := reflect.New(db.Statement.Schema.ModelType).Interface().(tenantScopedMarker)
|
|
return ok
|
|
}
|
|
|
|
// addTenantWhere：受租户模型 + ctx 有 tenant → 追加 tenant_id 过滤。
|
|
// ctx 无 tenant（系统/回填/未登录）→ 不过滤（这些路径本就需跨租户；用户面由中间件保证有 tenant）。
|
|
func addTenantWhere(db *gorm.DB) {
|
|
if !isTenantScopedStmt(db) {
|
|
return
|
|
}
|
|
if tid := tenantFromCtx(db.Statement.Context); tid != "" {
|
|
db.Statement.AddClause(clause.Where{Exprs: []clause.Expression{
|
|
clause.Eq{Column: clause.Column{Table: db.Statement.Table, Name: "tenant_id"}, Value: tid},
|
|
}})
|
|
}
|
|
}
|
|
|
|
// setTenantOnCreate：受租户模型 + ctx 有 tenant → 强制把 tenant_id 设为 ctx 租户（防越权写他租）。
|
|
func setTenantOnCreate(db *gorm.DB) {
|
|
if !isTenantScopedStmt(db) {
|
|
return
|
|
}
|
|
if tid := tenantFromCtx(db.Statement.Context); tid != "" {
|
|
db.Statement.SetColumn("TenantID", tid)
|
|
}
|
|
}
|
|
|
|
// BackfillRowTenants 给受租户表的存量行回填 tenant_id = owner 的默认租户（幂等：只填空的）。
|
|
// 系统操作、无 ctx 租户 → 不受插件过滤影响。当前覆盖 KB / Agent（同步路径已接插件）。
|
|
func (p *Postgres) BackfillRowTenants(ctx context.Context) error {
|
|
if p.db == nil {
|
|
return errStoreDisabled
|
|
}
|
|
// 子查询：owner(user) → 其默认(最早 active)租户。
|
|
sub := "SELECT tenant_id FROM sundynix_tenant_member m WHERE m.user_id = t.owner AND m.status = 'active' ORDER BY m.created_at ASC LIMIT 1"
|
|
for _, tbl := range []string{"sundynix_kb", "sundynix_agent"} {
|
|
q := "UPDATE " + tbl + " t SET tenant_id = (" + sub + ") WHERE (t.tenant_id IS NULL OR t.tenant_id = '')"
|
|
if err := p.db.WithContext(ctx).Exec(q).Error; err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|