diff --git a/docs/openapi/openapi.json b/docs/openapi/openapi.json index da486a8e81..d6682aefee 100644 --- a/docs/openapi/openapi.json +++ b/docs/openapi/openapi.json @@ -28815,6 +28815,10 @@ "type": "integer", "description": "Number of members in the team" }, + "virtual_key_count": { + "type": "integer", + "description": "Number of virtual keys assigned to the team" + }, "created_at": { "type": "string", "format": "date-time" @@ -28977,6 +28981,10 @@ "type": "integer", "description": "Number of members in the team" }, + "virtual_key_count": { + "type": "integer", + "description": "Number of virtual keys assigned to the team" + }, "created_at": { "type": "string", "format": "date-time" @@ -41525,6 +41533,12 @@ "in": "header", "name": "x-api-key", "description": "API key authentication via the `x-api-key` header.\nVirtual keys (prefixed with `sk-bf-`) can also be passed here.\n" + }, + "GoogleApiKeyAuth": { + "type": "apiKey", + "in": "header", + "name": "x-goog-api-key", + "description": "Google API key authentication via the `x-goog-api-key` header.\nVirtual keys (prefixed with `sk-bf-`) can also be passed here.\n" } }, "parameters": { @@ -53765,6 +53779,8 @@ }, "virtual_keys": { "type": "array", + "nullable": true, + "description": "Virtual keys assigned to this team. This field may be omitted or returned as null in some responses (for example, when a team is embedded inside a virtual-key response) to avoid nested `virtual_keys` recursion.\n", "items": { "$ref": "#/components/schemas/VirtualKey" } diff --git a/docs/openapi/schemas/management/governance.yaml b/docs/openapi/schemas/management/governance.yaml index bd04aecee7..f057682620 100644 --- a/docs/openapi/schemas/management/governance.yaml +++ b/docs/openapi/schemas/management/governance.yaml @@ -376,6 +376,11 @@ Team: $ref: '#/Budget' virtual_keys: type: array + nullable: true + description: > + Virtual keys assigned to this team. This field may be omitted or returned as null in some + responses (for example, when a team is embedded inside a virtual-key response) to avoid + nested `virtual_keys` recursion. items: $ref: '#/VirtualKey' profile: diff --git a/docs/openapi/schemas/management/users.yaml b/docs/openapi/schemas/management/users.yaml index 46db148f3e..aa9cb504d1 100644 --- a/docs/openapi/schemas/management/users.yaml +++ b/docs/openapi/schemas/management/users.yaml @@ -155,6 +155,9 @@ TeamObject: member_count: type: integer description: Number of members in the team + virtual_key_count: + type: integer + description: Number of virtual keys assigned to the team created_at: type: string format: date-time diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index 3dccd9eb99..e6f5650c3e 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -2519,10 +2519,14 @@ func (s *RDBConfigStore) DeleteVirtualKeyMCPConfig(ctx context.Context, id uint, return txDB.WithContext(ctx).Delete(&tables.TableVirtualKeyMCPConfig{}, "id = ?", id).Error } +const teamSelectWithVKCount = "governance_teams.*, (SELECT COUNT(*) FROM governance_virtual_keys WHERE team_id = governance_teams.id) AS virtual_key_count" + // GetTeams retrieves all teams from the database. func (s *RDBConfigStore) GetTeams(ctx context.Context, customerID string) ([]tables.TableTeam, error) { // Preload relationships for complete information - query := s.db.WithContext(ctx).Preload("Customer").Preload("Budget").Preload("RateLimit") + query := s.db.WithContext(ctx). + Select(teamSelectWithVKCount). + Preload("Customer").Preload("Budget").Preload("RateLimit") // Optional filtering by customer if customerID != "" { query = query.Where("customer_id = ?", customerID) @@ -2564,6 +2568,7 @@ func (s *RDBConfigStore) GetTeamsPaginated(ctx context.Context, params TeamsQuer var teams []tables.TableTeam if err := baseQuery. + Select(teamSelectWithVKCount). Preload("Customer").Preload("Budget").Preload("RateLimit"). Order("created_at ASC, id ASC"). Offset(offset).Limit(limit). @@ -2577,7 +2582,10 @@ func (s *RDBConfigStore) GetTeamsPaginated(ctx context.Context, params TeamsQuer // GetTeam retrieves a specific team from the database. func (s *RDBConfigStore) GetTeam(ctx context.Context, id string) (*tables.TableTeam, error) { var team tables.TableTeam - if err := s.db.WithContext(ctx).Preload("Customer").Preload("Budget").Preload("RateLimit").First(&team, "id = ?", id).Error; err != nil { + if err := s.db.WithContext(ctx). + Select(teamSelectWithVKCount). + Preload("Customer").Preload("Budget").Preload("RateLimit"). + First(&team, "id = ?", id).Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, ErrNotFound } @@ -3425,7 +3433,9 @@ func (s *RDBConfigStore) GetGovernanceConfig(ctx context.Context) (*GovernanceCo Find(&virtualKeys).Error; err != nil { return nil, err } - if err := s.db.WithContext(ctx).Find(&teams).Error; err != nil { + if err := s.db.WithContext(ctx). + Select(teamSelectWithVKCount). + Find(&teams).Error; err != nil { return nil, err } if err := s.db.WithContext(ctx).Find(&customers).Error; err != nil { @@ -4491,4 +4501,4 @@ func (s *RDBConfigStore) TransferOauthUserTokensFromGatewaySession(ctx context.C } s.logger.Debug("[rdb] TransferOauthUserTokensFromGatewaySession done: rows_affected=%d", result.RowsAffected) return nil -} \ No newline at end of file +} diff --git a/framework/configstore/tables/team.go b/framework/configstore/tables/team.go index 3054263ce8..e96614c600 100644 --- a/framework/configstore/tables/team.go +++ b/framework/configstore/tables/team.go @@ -20,7 +20,10 @@ type TableTeam struct { Customer *TableCustomer `gorm:"foreignKey:CustomerID" json:"customer,omitempty"` Budget *TableBudget `gorm:"foreignKey:BudgetID" json:"budget,omitempty"` RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID" json:"rate_limit,omitempty"` - VirtualKeys []TableVirtualKey `gorm:"foreignKey:TeamID" json:"virtual_keys"` + VirtualKeys []TableVirtualKey `gorm:"foreignKey:TeamID" json:"virtual_keys,omitempty"` + + // Computed (not a DB column) — populated via correlated subquery in query layer, hence no migration + VirtualKeyCount int64 `gorm:"->;-:migration" json:"virtual_key_count"` Profile *string `gorm:"type:text" json:"-"` ParsedProfile map[string]interface{} `gorm:"-" json:"profile"` diff --git a/plugins/governance/store.go b/plugins/governance/store.go index 4775c71bd4..ed4e3bce61 100644 --- a/plugins/governance/store.go +++ b/plugins/governance/store.go @@ -283,6 +283,9 @@ func (gs *LocalGovernanceStore) GetGovernanceData() *GovernanceData { } clone := *team refreshTeamAssociations(&clone) + // Reset to 0 — will be recomputed from live VKs below to stay accurate + // after creates/updates/deletes that don't trigger a full ReloadTeam. + clone.VirtualKeyCount = 0 teams[key.(string)] = &clone return true // continue iteration }) @@ -313,21 +316,6 @@ func (gs *LocalGovernanceStore) GetGovernanceData() *GovernanceData { return true // continue iteration }) - for _, team := range teams { - if team == nil { - continue - } - if team.CustomerID != nil { - if customer, exists := customers[*team.CustomerID]; exists && customer != nil { - team.Customer = customer - - nestedTeam := *team - nestedTeam.Customer = nil - customer.Teams = append(customer.Teams, nestedTeam) - } - } - } - for _, vk := range virtualKeys { if vk == nil { continue @@ -335,6 +323,7 @@ func (gs *LocalGovernanceStore) GetGovernanceData() *GovernanceData { if vk.TeamID != nil { if team, exists := teams[*vk.TeamID]; exists && team != nil { vk.Team = team + team.VirtualKeyCount++ } } if vk.CustomerID != nil { @@ -348,6 +337,21 @@ func (gs *LocalGovernanceStore) GetGovernanceData() *GovernanceData { } } + for _, team := range teams { + if team == nil { + continue + } + if team.CustomerID != nil { + if customer, exists := customers[*team.CustomerID]; exists && customer != nil { + team.Customer = customer + + nestedTeam := *team + nestedTeam.Customer = nil + customer.Teams = append(customer.Teams, nestedTeam) + } + } + } + for _, customer := range customers { if customer == nil { continue