From e4f94676a0668c46f07a280314c3eb9d557efa21 Mon Sep 17 00:00:00 2001 From: Anuj Parihar Date: Thu, 11 Jun 2026 13:17:07 +0530 Subject: [PATCH] feat: add unique constraint migration on customer table name field --- framework/configstore/migrations.go | 87 ++++++++++++++++++++++++ framework/configstore/tables/customer.go | 4 +- transports/bifrost-http/lib/config.go | 28 ++++++-- 3 files changed, 112 insertions(+), 7 deletions(-) diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 3a849a5c7cb..720930a0e0c 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -877,6 +877,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationAddFastModePricingColumns(ctx, db); err != nil { return err } + if err := migrationAddCustomerNameUniqueConstraint(ctx, db); err != nil { + return err + } return nil } @@ -9962,3 +9965,87 @@ func migrationAddMCPLibrarySourceColumns(ctx context.Context, db *gorm.DB) error } return nil } + +// migrationAddCustomerNameUniqueConstraint deduplicates governance_customers by +// appending -1, -2, … to later occurrences of the same name (ordered by +// created_at then id), then adds a unique index on the name column. +func migrationAddCustomerNameUniqueConstraint(ctx context.Context, db *gorm.DB) error { + const idxName = "idx_governance_customers_name" + + // Step 1 (transactional): rename duplicate customer names so the later + // CREATE UNIQUE INDEX cannot fail due to pre-existing duplicates. + if err := RunSingleMigration(ctx, nil, db, &migrator.Migration{ + ID: "add_customer_name_unique_constraint_dedup", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + + // Fetch all customers in a stable order so the earliest-created row + // always keeps the original name and later duplicates receive suffixes. + var customers []tables.TableCustomer + if err := tx.Order("created_at ASC, id ASC").Find(&customers).Error; err != nil { + return fmt.Errorf("failed to fetch customers: %w", err) + } + + // taken tracks every name that is currently (or will be) in use so + // suffix search never collides with an existing original name. + taken := make(map[string]bool, len(customers)) + for _, c := range customers { + taken[c.Name] = true + } + + firstSeen := make(map[string]bool, len(customers)) + for _, c := range customers { + if !firstSeen[c.Name] { + firstSeen[c.Name] = true + continue // earliest occurrence keeps its name + } + // Find the lowest suffix whose candidate name is not already taken. + suffix := 1 + candidate := fmt.Sprintf("%s-%d", c.Name, suffix) + for taken[candidate] { + suffix++ + candidate = fmt.Sprintf("%s-%d", c.Name, suffix) + } + taken[candidate] = true + if err := tx.Model(&tables.TableCustomer{}).Where("id = ?", c.ID).Update("name", candidate).Error; err != nil { + return fmt.Errorf("failed to rename customer %s to %q: %w", c.ID, candidate, err) + } + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + return nil // name renames are not reversed; dropping the index in step 2 restores the invariant + }, + }); err != nil { + return err + } + + // Step 2 (non-transactional): create the unique index. + // UseTransaction must be false because CREATE INDEX CONCURRENTLY cannot + // execute inside a transaction block. IF NOT EXISTS makes this step safe + // to re-run if the process crashes after the index is built but before + // the migration record is written. + noTxOpts := *migrator.DefaultOptions + noTxOpts.UseTransaction = false + return RunSingleMigration(ctx, &noTxOpts, db, &migrator.Migration{ + ID: "add_customer_name_unique_constraint_index", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + // SQLite does not support CONCURRENTLY; use the plain form there. + var stmt string + if tx.Dialector.Name() == "sqlite" { + stmt = "CREATE UNIQUE INDEX IF NOT EXISTS " + idxName + " ON governance_customers (name)" + } else { + stmt = "CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS " + idxName + " ON governance_customers (name)" + } + if err := tx.Exec(stmt).Error; err != nil { + return fmt.Errorf("failed to create unique index on governance_customers.name: %w", err) + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + return tx.Exec("DROP INDEX IF EXISTS " + idxName).Error + }, + }) +} diff --git a/framework/configstore/tables/customer.go b/framework/configstore/tables/customer.go index f70574f59a5..bb8a3a2e95d 100644 --- a/framework/configstore/tables/customer.go +++ b/framework/configstore/tables/customer.go @@ -9,14 +9,14 @@ import ( // TableCustomer represents a customer entity with budgets, rate limit and team/VK association type TableCustomer struct { ID string `gorm:"primaryKey;type:varchar(255)" json:"id"` - Name string `gorm:"type:varchar(255);not null" json:"name"` + Name string `gorm:"type:varchar(255);not null;uniqueIndex:idx_governance_customers_name" json:"name"` RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"` // BudgetID is a config-file-only field referencing a pre-declared budget (from governance.budgets) to link to this customer. Not persisted; used by the config sync path to set customer_id on the referenced budget row. BudgetID *string `gorm:"-" json:"budget_id,omitempty"` // Relationships - Budgets []TableBudget `gorm:"foreignKey:CustomerID;constraint:OnDelete:CASCADE" json:"budgets,omitempty"` + Budgets []TableBudget `gorm:"foreignKey:CustomerID;constraint:OnDelete:CASCADE" json:"budgets,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"` diff --git a/transports/bifrost-http/lib/config.go b/transports/bifrost-http/lib/config.go index cb2b71e8d15..351c36dfe2d 100644 --- a/transports/bifrost-http/lib/config.go +++ b/transports/bifrost-http/lib/config.go @@ -2066,7 +2066,13 @@ func mergeGovernanceConfig(ctx context.Context, config *Config, configData *Conf found := false for j, existingCustomer := range governanceConfig.Customers { - if existingCustomer.ID == newCustomer.ID { + idMatch := existingCustomer.ID == newCustomer.ID + nameMatch := newCustomer.ID == "" && existingCustomer.Name == newCustomer.Name + if idMatch || nameMatch { + if nameMatch { + // Config file has no ID; adopt the DB record's ID so updates use the right primary key. + configData.Governance.Customers[i].ID = existingCustomer.ID + } found = true if existingCustomer.ConfigHash != fileCustomerHash { logger.Debug("config hash mismatch for customer %s, syncing from config file", newCustomer.ID) @@ -2081,6 +2087,9 @@ func mergeGovernanceConfig(ctx context.Context, config *Config, configData *Conf } if !found { configData.Governance.Customers[i].ConfigHash = fileCustomerHash + if configData.Governance.Customers[i].ID == "" { + configData.Governance.Customers[i].ID = uuid.NewString() + } customersToAdd = append(customersToAdd, configData.Governance.Customers[i]) } } @@ -2097,7 +2106,13 @@ func mergeGovernanceConfig(ctx context.Context, config *Config, configData *Conf found := false for j, existingTeam := range governanceConfig.Teams { - if existingTeam.ID == newTeam.ID { + idMatch := existingTeam.ID == newTeam.ID + nameMatch := newTeam.ID == "" && existingTeam.Name == newTeam.Name + if idMatch || nameMatch { + if nameMatch { + // Config file has no ID; adopt the DB record's ID so updates use the right primary key. + configData.Governance.Teams[i].ID = existingTeam.ID + } found = true if existingTeam.ConfigHash != fileTeamHash { logger.Debug("config hash mismatch for team %s, syncing from config file", newTeam.ID) @@ -2112,6 +2127,9 @@ func mergeGovernanceConfig(ctx context.Context, config *Config, configData *Conf } if !found { configData.Governance.Teams[i].ConfigHash = fileTeamHash + if configData.Governance.Teams[i].ID == "" { + configData.Governance.Teams[i].ID = uuid.NewString() + } teamsToAdd = append(teamsToAdd, configData.Governance.Teams[i]) } } @@ -2173,6 +2191,9 @@ func mergeGovernanceConfig(ctx context.Context, config *Config, configData *Conf } if !found { configData.Governance.VirtualKeys[i].ConfigHash = fileVKHash + if configData.Governance.VirtualKeys[i].ID == "" { + configData.Governance.VirtualKeys[i].ID = uuid.NewString() + } // if the virtual key value is env.VIRTUAL_KEY_VALUE, then we will need to resolve the environment variable // Process environment variable for virtual key value if strings.HasPrefix(configData.Governance.VirtualKeys[i].Value, "env.") { @@ -2779,9 +2800,6 @@ func updateGovernanceConfigInStore( // Create virtual keys with explicit association handling for i := range virtualKeysToAdd { virtualKey := &virtualKeysToAdd[i] - if virtualKey.ID == "" { - virtualKey.ID = uuid.NewString() - } providerConfigs := virtualKey.ProviderConfigs mcpConfigs := virtualKey.MCPConfigs virtualKey.ProviderConfigs = nil