package store import ( "context" "gorm.io/gorm" "gorm.io/gorm/clause" ) // 角色/成员状态常量。 const ( RoleOwner = "owner" RoleAdmin = "admin" RoleMember = "member" RoleViewer = "viewer" RoleBillingAdmin = "billing_admin" ) // CreateTenant 建一个租户。 func (p *Postgres) CreateTenant(ctx context.Context, name, slug, plan string) (*Tenant, error) { if p.db == nil { return nil, errStoreDisabled } if plan == "" { plan = "free" } t := &Tenant{Name: name, Slug: slug, Plan: plan, Status: "active"} if err := p.db.WithContext(ctx).Create(t).Error; err != nil { return nil, err } return t, nil } // AddMember 添加成员(幂等:已存在则不重复插入)。 func (p *Postgres) AddMember(ctx context.Context, tenantID, userID, role string) error { if p.db == nil { return errStoreDisabled } m := &TenantMember{TenantID: tenantID, UserID: userID, Role: role, Status: "active"} // 命中唯一约束(tenant_id,user_id)则忽略,保证幂等。 return p.db.WithContext(ctx). Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "tenant_id"}, {Name: "user_id"}}, DoNothing: true}). Create(m).Error } // DefaultTenantForUser 返回用户所属的默认租户(当前取其最早加入的 active 成员关系;无则 nil)。 // tenantIDForOwner 返回某 user 的默认(最早 active)租户 id;无则空串。 // 供异步路径(入库 / 评测回写,无请求 ctx 租户)在写入时补 tenant_id。 func (p *Postgres) tenantIDForOwner(ctx context.Context, owner string) string { if p.db == nil || owner == "" { return "" } var m TenantMember if err := p.db.WithContext(ctx).Select("tenant_id"). Where("user_id = ? AND status = ?", owner, "active"). Order("created_at asc").First(&m).Error; err != nil { return "" } return m.TenantID } func (p *Postgres) DefaultTenantForUser(ctx context.Context, userID string) (*Tenant, error) { if p.db == nil { return nil, errStoreDisabled } var m TenantMember err := p.db.WithContext(ctx). Where("user_id = ? AND status = ?", userID, "active"). Order("created_at asc").First(&m).Error if err != nil { if err == gorm.ErrRecordNotFound { return nil, nil } return nil, err } var t Tenant if err := p.db.WithContext(ctx).First(&t, "id = ?", m.TenantID).Error; err != nil { return nil, err } return &t, nil } // EnsureDefaultTenant 幂等地保证用户有一个默认租户:已有 active 成员关系则返回其默认租户; // 否则建一个单人租户(owner) 并返回。用于注册与存量回填。 func (p *Postgres) EnsureDefaultTenant(ctx context.Context, userID, name string) (*Tenant, error) { if p.db == nil { return nil, errStoreDisabled } if t, err := p.DefaultTenantForUser(ctx, userID); err != nil { return nil, err } else if t != nil { return t, nil } if name == "" { name = "个人工作区" } t, err := p.CreateTenant(ctx, name, "default-"+userID, "free") if err != nil { return nil, err } if err := p.AddMember(ctx, t.ID, userID, RoleOwner); err != nil { return nil, err } return t, nil } // BackfillDefaultTenants 给所有还没有租户成员关系的存量用户补建默认租户(启动时幂等调用)。 func (p *Postgres) BackfillDefaultTenants(ctx context.Context) (int, error) { if p.db == nil { return 0, errStoreDisabled } var users []User // 无 active 成员关系的用户。 if err := p.db.WithContext(ctx). Where("id NOT IN (?)", p.db.Model(&TenantMember{}).Select("user_id").Where("status = ?", "active")). Find(&users).Error; err != nil { return 0, err } n := 0 for _, u := range users { if _, err := p.EnsureDefaultTenant(ctx, u.ID, firstNonEmpty(u.Name, u.Email)); err != nil { return n, err } n++ } return n, nil } // GetTenant 按 id 取租户(不存在返回 nil)。 func (p *Postgres) GetTenant(ctx context.Context, id string) (*Tenant, error) { if p.db == nil { return nil, errStoreDisabled } var t Tenant if err := p.db.WithContext(ctx).First(&t, "id = ?", id).Error; err != nil { if err == gorm.ErrRecordNotFound { return nil, nil } return nil, err } return &t, nil } // MemberRole 返回用户在某租户的角色(无成员关系返回空)。 func (p *Postgres) MemberRole(ctx context.Context, tenantID, userID string) string { if p.db == nil { return "" } var m TenantMember if err := p.db.WithContext(ctx). Where("tenant_id = ? AND user_id = ? AND status = ?", tenantID, userID, "active"). First(&m).Error; err != nil { return "" } return m.Role } func firstNonEmpty(a, b string) string { if a != "" { return a } return b }