Skip to content

feat: adds rate limits to customers and teams #1619

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
akshaydeo merged 1 commit into from
Feb 12, 2026
Merged

feat: adds rate limits to customers and teams #1619

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions framework/configstore/migrations.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -247,6 +247,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
if err := migrationAddReplicateDeploymentsJSONColumn(ctx, db); err != nil {
return err
}
if err := migrationAddRateLimitToTeamsAndCustomers(ctx, db); err != nil {
return err
}
return nil
}

Expand Down Expand Up @@ -3357,3 +3360,52 @@ func migrationAddReplicateDeploymentsJSONColumn(ctx context.Context, db *gorm.DB
}
return nil
}

// migrationAddRateLimitToTeamsAndCustomers adds rate_limit_id column to governance_teams and governance_customers tables
func migrationAddRateLimitToTeamsAndCustomers(ctx context.Context, db *gorm.DB) error {
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
ID: "add_rate_limit_to_teams_and_customers",
Migrate: func(tx *gorm.DB) error {
tx = tx.WithContext(ctx)
migrator := tx.Migrator()

// Add rate_limit_id to governance_teams table
if !migrator.HasColumn(&tables.TableTeam{}, "rate_limit_id") {
if err := migrator.AddColumn(&tables.TableTeam{}, "rate_limit_id"); err != nil {
return fmt.Errorf("failed to add rate_limit_id column to teams: %w", err)
}
}

// Add rate_limit_id to governance_customers table
if !migrator.HasColumn(&tables.TableCustomer{}, "rate_limit_id") {
if err := migrator.AddColumn(&tables.TableCustomer{}, "rate_limit_id"); err != nil {
return fmt.Errorf("failed to add rate_limit_id column to customers: %w", err)
}
}

return nil
},
Rollback: func(tx *gorm.DB) error {
tx = tx.WithContext(ctx)
migrator := tx.Migrator()

if migrator.HasColumn(&tables.TableTeam{}, "rate_limit_id") {
if err := migrator.DropColumn(&tables.TableTeam{}, "rate_limit_id"); err != nil {
return fmt.Errorf("failed to drop rate_limit_id column from teams: %w", err)
}
}

if migrator.HasColumn(&tables.TableCustomer{}, "rate_limit_id") {
if err := migrator.DropColumn(&tables.TableCustomer{}, "rate_limit_id"); err != nil {
return fmt.Errorf("failed to drop rate_limit_id column from customers: %w", err)
}
}

return nil
},
}})
if err := m.Migrate(); err != nil {
return fmt.Errorf("error running rate limit migration for teams and customers: %s", err.Error())
}
return nil
}
68 changes: 58 additions & 10 deletions framework/configstore/rdb.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1812,7 +1812,7 @@ func (s *RDBConfigStore) DeleteVirtualKeyMCPConfig(ctx context.Context, id uint,
// 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")
query := s.db.WithContext(ctx).Preload("Customer").Preload("Budget").Preload("RateLimit")
// Optional filtering by customer
if customerID != "" {
query = query.Where("customer_id = ?", customerID)
Expand All @@ -1827,7 +1827,7 @@ func (s *RDBConfigStore) GetTeams(ctx context.Context, customerID string) ([]tab
// 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").First(&team, "id = ?", id).Error; err != nil {
if err := s.db.WithContext(ctx).Preload("Customer").Preload("Budget").Preload("RateLimit").First(&team, "id = ?", id).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, ErrNotFound
}
Expand Down Expand Up @@ -1868,7 +1868,7 @@ func (s *RDBConfigStore) UpdateTeam(ctx context.Context, team *tables.TableTeam,
func (s *RDBConfigStore) DeleteTeam(ctx context.Context, id string) error {
if err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
var team tables.TableTeam
if err := tx.WithContext(ctx).Preload("Budget").First(&team, "id = ?", id).Error; err != nil {
if err := tx.WithContext(ctx).Preload("Budget").Preload("RateLimit").First(&team, "id = ?", id).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return ErrNotFound
}
Expand All @@ -1878,8 +1878,9 @@ func (s *RDBConfigStore) DeleteTeam(ctx context.Context, id string) error {
if err := tx.WithContext(ctx).Model(&tables.TableVirtualKey{}).Where("team_id = ?", id).Update("team_id", nil).Error; err != nil {
return err
}
// Store the budget ID before deleting the team
// Store the budget and rate limit IDs before deleting the team
budgetID := team.BudgetID
rateLimitID := team.RateLimitID
// Delete the team first
if err := tx.WithContext(ctx).Delete(&tables.TableTeam{}, "id = ?", id).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
Expand All @@ -1893,6 +1894,12 @@ func (s *RDBConfigStore) DeleteTeam(ctx context.Context, id string) error {
return err
}
}
// Delete the team's rate limit if it exists
if rateLimitID != nil {
if err := tx.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *rateLimitID).Error; err != nil {
return err
}
}
return nil
}); err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
Expand All @@ -1906,7 +1913,7 @@ func (s *RDBConfigStore) DeleteTeam(ctx context.Context, id string) error {
// GetCustomers retrieves all customers from the database.
func (s *RDBConfigStore) GetCustomers(ctx context.Context) ([]tables.TableCustomer, error) {
var customers []tables.TableCustomer
if err := s.db.WithContext(ctx).Preload("Teams").Preload("Budget").Order("created_at ASC").Find(&customers).Error; err != nil {
if err := s.db.WithContext(ctx).Preload("Teams").Preload("Budget").Preload("RateLimit").Order("created_at ASC").Find(&customers).Error; err != nil {
return nil, err
}
return customers, nil
Expand All @@ -1915,7 +1922,7 @@ func (s *RDBConfigStore) GetCustomers(ctx context.Context) ([]tables.TableCustom
// GetCustomer retrieves a specific customer from the database.
func (s *RDBConfigStore) GetCustomer(ctx context.Context, id string) (*tables.TableCustomer, error) {
var customer tables.TableCustomer
if err := s.db.WithContext(ctx).Preload("Teams").Preload("Budget").First(&customer, "id = ?", id).Error; err != nil {
if err := s.db.WithContext(ctx).Preload("Teams").Preload("Budget").Preload("RateLimit").First(&customer, "id = ?", id).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, ErrNotFound
}
Expand Down Expand Up @@ -1956,7 +1963,7 @@ func (s *RDBConfigStore) UpdateCustomer(ctx context.Context, customer *tables.Ta
func (s *RDBConfigStore) DeleteCustomer(ctx context.Context, id string) error {
if err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
var customer tables.TableCustomer
if err := tx.WithContext(ctx).Preload("Budget").First(&customer, "id = ?", id).Error; err != nil {
if err := tx.WithContext(ctx).Preload("Budget").Preload("RateLimit").First(&customer, "id = ?", id).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return ErrNotFound
}
Expand All @@ -1970,8 +1977,9 @@ func (s *RDBConfigStore) DeleteCustomer(ctx context.Context, id string) error {
if err := tx.WithContext(ctx).Model(&tables.TableTeam{}).Where("customer_id = ?", id).Update("customer_id", nil).Error; err != nil {
return err
}
// Store the budget ID before deleting the customer
// Store the budget and rate limit IDs before deleting the customer
budgetID := customer.BudgetID
rateLimitID := customer.RateLimitID
// Delete the customer first
if err := tx.WithContext(ctx).Delete(&tables.TableCustomer{}, "id = ?", id).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
Expand All @@ -1985,6 +1993,12 @@ func (s *RDBConfigStore) DeleteCustomer(ctx context.Context, id string) error {
return err
}
}
// Delete the customer's rate limit if it exists
if rateLimitID != nil {
if err := tx.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", *rateLimitID).Error; err != nil {
return err
}
}
return nil
}); err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
Expand All @@ -2005,9 +2019,15 @@ func (s *RDBConfigStore) GetRateLimits(ctx context.Context) ([]tables.TableRateL
}

// GetRateLimit retrieves a specific rate limit from the database.
func (s *RDBConfigStore) GetRateLimit(ctx context.Context, id string) (*tables.TableRateLimit, error) {
func (s *RDBConfigStore) GetRateLimit(ctx context.Context, id string, tx ...*gorm.DB) (*tables.TableRateLimit, error) {
var txDB *gorm.DB
if len(tx) > 0 {
txDB = tx[0]
} else {
txDB = s.db
}
var rateLimit tables.TableRateLimit
if err := s.db.WithContext(ctx).First(&rateLimit, "id = ?", id).Error; err != nil {
if err := txDB.WithContext(ctx).First(&rateLimit, "id = ?", id).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, ErrNotFound
}
Expand Down Expand Up @@ -2060,6 +2080,20 @@ func (s *RDBConfigStore) UpdateRateLimits(ctx context.Context, rateLimits []*tab
return nil
}

// DeleteRateLimit deletes a rate limit from the database.
func (s *RDBConfigStore) DeleteRateLimit(ctx context.Context, id string, tx ...*gorm.DB) error {
var txDB *gorm.DB
if len(tx) > 0 {
txDB = tx[0]
} else {
txDB = s.db
}
if err := txDB.WithContext(ctx).Delete(&tables.TableRateLimit{}, "id = ?", id).Error; err != nil {
return s.parseGormError(err)
}
return nil
}

// GetBudgets retrieves all budgets from the database.
func (s *RDBConfigStore) GetBudgets(ctx context.Context) ([]tables.TableBudget, error) {
var budgets []tables.TableBudget
Expand Down Expand Up @@ -2131,6 +2165,20 @@ func (s *RDBConfigStore) UpdateBudget(ctx context.Context, budget *tables.TableB
return nil
}

// DeleteBudget deletes a budget from the database.
func (s *RDBConfigStore) DeleteBudget(ctx context.Context, id string, tx ...*gorm.DB) error {
var txDB *gorm.DB
if len(tx) > 0 {
txDB = tx[0]
} else {
txDB = s.db
}
if err := txDB.WithContext(ctx).Delete(&tables.TableBudget{}, "id = ?", id).Error; err != nil {
return s.parseGormError(err)
}
return nil
}

// UpdateBudgetUsage updates only the current_usage field of a budget.
// Uses SkipHooks to avoid triggering BeforeSave validation since we're only updating usage.
func (s *RDBConfigStore) UpdateBudgetUsage(ctx context.Context, id string, currentUsage float64) error {
Expand Down
4 changes: 3 additions & 1 deletion framework/configstore/store.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,17 +102,19 @@ type ConfigStore interface {

// Rate limit CRUD
GetRateLimits(ctx context.Context) ([]tables.TableRateLimit, error)
GetRateLimit(ctx context.Context, id string) (*tables.TableRateLimit, error)
GetRateLimit(ctx context.Context, id string, tx ...*gorm.DB) (*tables.TableRateLimit, error)
CreateRateLimit(ctx context.Context, rateLimit *tables.TableRateLimit, tx ...*gorm.DB) error
UpdateRateLimit(ctx context.Context, rateLimit *tables.TableRateLimit, tx ...*gorm.DB) error
UpdateRateLimits(ctx context.Context, rateLimits []*tables.TableRateLimit, tx ...*gorm.DB) error
DeleteRateLimit(ctx context.Context, id string, tx ...*gorm.DB) error

Comment thread
coderabbitai[bot] marked this conversation as resolved.
// Budget CRUD
GetBudgets(ctx context.Context) ([]tables.TableBudget, error)
GetBudget(ctx context.Context, id string, tx ...*gorm.DB) (*tables.TableBudget, error)
CreateBudget(ctx context.Context, budget *tables.TableBudget, tx ...*gorm.DB) error
UpdateBudget(ctx context.Context, budget *tables.TableBudget, tx ...*gorm.DB) error
UpdateBudgets(ctx context.Context, budgets []*tables.TableBudget, tx ...*gorm.DB) error
DeleteBudget(ctx context.Context, id string, tx ...*gorm.DB) error
UpdateBudgetUsage(ctx context.Context, id string, currentUsage float64) error
UpdateRateLimitUsage(ctx context.Context, id string, tokenCurrentUsage int64, requestCurrentUsage int64) error

Expand Down
10 changes: 6 additions & 4 deletions framework/configstore/tables/customer.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,16 @@ package tables

import "time"

// TableCustomer represents a customer entity with budget
// TableCustomer represents a customer entity with budget and rate limit
type TableCustomer struct {
ID string `gorm:"primaryKey;type:varchar(255)" json:"id"`
Name string `gorm:"type:varchar(255);not null" json:"name"`
BudgetID *string `gorm:"type:varchar(255);index" json:"budget_id,omitempty"`
ID string `gorm:"primaryKey;type:varchar(255)" json:"id"`
Name string `gorm:"type:varchar(255);not null" json:"name"`
BudgetID *string `gorm:"type:varchar(255);index" json:"budget_id,omitempty"`
RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"`

// Relationships
Budget *TableBudget `gorm:"foreignKey:BudgetID" json:"budget,omitempty"`
RateLimit *TableRateLimit `gorm:"foreignKey:RateLimitID" json:"rate_limit,omitempty"`
Teams []TableTeam `gorm:"foreignKey:CustomerID" json:"teams"`
VirtualKeys []TableVirtualKey `gorm:"foreignKey:CustomerID" json:"virtual_keys"`

Expand Down
12 changes: 7 additions & 5 deletions framework/configstore/tables/team.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,16 +8,18 @@ import (
"gorm.io/gorm"
)

// TableTeam represents a team entity with budget and customer association
// TableTeam represents a team entity with budget, rate limit and customer association
type TableTeam struct {
ID string `gorm:"primaryKey;type:varchar(255)" json:"id"`
Name string `gorm:"type:varchar(255);not null" json:"name"`
CustomerID *string `gorm:"type:varchar(255);index" json:"customer_id,omitempty"` // A team can belong to a customer
BudgetID *string `gorm:"type:varchar(255);index" json:"budget_id,omitempty"`
ID string `gorm:"primaryKey;type:varchar(255)" json:"id"`
Name string `gorm:"type:varchar(255);not null" json:"name"`
CustomerID *string `gorm:"type:varchar(255);index" json:"customer_id,omitempty"` // A team can belong to a customer
BudgetID *string `gorm:"type:varchar(255);index" json:"budget_id,omitempty"`
RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"`

// Relationships
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"`

Profile *string `gorm:"type:text" json:"-"`
Expand Down
Loading
Loading