From e33c29b38b9ef2dd4cf516624cf5b1bbd5a5ebf7 Mon Sep 17 00:00:00 2001 From: roroghost17 Date: Sat, 6 Jun 2026 12:49:53 +0530 Subject: [PATCH] fix: supports backup for governance migration --- framework/configstore/migrations.go | 86 +++++++++++++---- framework/configstore/migrations_test.go | 115 +++++++++++++++++++++++ 2 files changed, 182 insertions(+), 19 deletions(-) diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 4449946908c..ae9637974cc 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -188,7 +188,6 @@ type legacyBudgetTeam struct { // TableName returns the governance_teams table name for legacyBudgetTeam. func (legacyBudgetTeam) TableName() string { return "governance_teams" } - // sqliteColumnInfo holds the information about a SQLite column. type sqliteColumnInfo struct { Name string `gorm:"column:name"` @@ -3933,6 +3932,15 @@ func migrationAddModelConfigScopeColumns(ctx context.Context, db *gorm.DB) error // (scope='global', provider=, model_name='*') "all models on this provider" rows, // reusing the same budget/rate-limit rows. It then NULLs the provider FKs so the old // provider-governance enforcement path goes inert (single source of truth = model_configs). +// +// If a user-created (global, provider, '*') row already occupies the wildcard slot, +// the provider's own budget/rate-limit links are dropped by the detach without being +// folded anywhere. For exactly those providers the stranded rows are snapshotted into +// governance_provider_budgets_backup / governance_provider_rate_limits_backup +// (data-only copies keyed by provider_name) — the support-facing record from which +// that governance can be restored on request. The tables are only created when such +// a conflict actually occurs, and a backup failure is logged and skipped rather than +// failing the migration. func migrationMigrateProviderGovernanceToModelConfigs(ctx context.Context, db *gorm.DB) error { m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ ID: "migrate_provider_governance_to_model_configs", @@ -3987,24 +3995,64 @@ func migrationMigrateProviderGovernanceToModelConfigs(ctx context.Context, db *g }).Error; err != nil { return fmt.Errorf("failed to create wildcard model config for provider %q: %w", p.Name, err) } - case err != nil: - return fmt.Errorf("failed to check existing wildcard config for provider %q: %w", p.Name, err) - default: - // A wildcard row already exists: backfill only the governance slots it lacks, - // so we never overwrite values already present. This is commutative and safe - // to repeat (a clean re-run finds nothing to fill and writes nothing). - updates := map[string]any{} - if existing.BudgetID == nil && p.BudgetID != nil { - updates["budget_id"] = p.BudgetID - } - if existing.RateLimitID == nil && p.RateLimitID != nil { - updates["rate_limit_id"] = p.RateLimitID - } - if len(updates) > 0 { - updates["updated_at"] = now - if err := tx.Table((tables.TableModelConfig{}).TableName()). - Where("id = ?", existing.ID).Updates(updates).Error; err != nil { - return fmt.Errorf("failed to merge provider governance into wildcard model config for provider %q: %w", p.Name, err) + } else { + // Conflict: a user-created wildcard row already occupies this provider's + // (global, '*') slot, so the provider's own governance is not folded + // anywhere and the detach below would strand it. Snapshot the stranded + // rows into backup tables so support can restore them on request. + // + // A backup failure must never block the migration. The statements run + // inside a savepoint so a failed one can be rolled back without + // aborting the surrounding transaction, then logged and skipped. + const backupSavepoint = "sp_provider_gov_backup" + if spErr := tx.SavePoint(backupSavepoint).Error; spErr != nil { + log.Printf("[configstore] could not create savepoint for governance backup of provider %q (skipping backup): %v", p.Name, spErr) + } else if backupErr := func() error { + if p.BudgetID != nil { + if err := tx.Exec(` + CREATE TABLE IF NOT EXISTS governance_provider_budgets_backup AS + SELECT p.name AS provider_name, b.* + FROM config_providers p + INNER JOIN governance_budgets b ON b.id = p.budget_id + WHERE 1 = 0 + `).Error; err != nil { + return fmt.Errorf("failed to create provider budgets backup table: %w", err) + } + if err := tx.Exec(` + INSERT INTO governance_provider_budgets_backup + SELECT p.name AS provider_name, b.* + FROM config_providers p + INNER JOIN governance_budgets b ON b.id = p.budget_id + WHERE p.name = ? + `, p.Name).Error; err != nil { + return fmt.Errorf("failed to back up budget: %w", err) + } + } + if p.RateLimitID != nil { + if err := tx.Exec(` + CREATE TABLE IF NOT EXISTS governance_provider_rate_limits_backup AS + SELECT p.name AS provider_name, rl.* + FROM config_providers p + INNER JOIN governance_rate_limits rl ON rl.id = p.rate_limit_id + WHERE 1 = 0 + `).Error; err != nil { + return fmt.Errorf("failed to create provider rate limits backup table: %w", err) + } + if err := tx.Exec(` + INSERT INTO governance_provider_rate_limits_backup + SELECT p.name AS provider_name, rl.* + FROM config_providers p + INNER JOIN governance_rate_limits rl ON rl.id = p.rate_limit_id + WHERE p.name = ? + `, p.Name).Error; err != nil { + return fmt.Errorf("failed to back up rate limit: %w", err) + } + } + return nil + }(); backupErr != nil { + log.Printf("[configstore] failed to back up stranded governance for provider %q (continuing without backup): %v", p.Name, backupErr) + if rbErr := tx.RollbackTo(backupSavepoint).Error; rbErr != nil { + log.Printf("[configstore] could not roll back governance backup savepoint for provider %q: %v", p.Name, rbErr) } } } diff --git a/framework/configstore/migrations_test.go b/framework/configstore/migrations_test.go index f522c7cb955..986303400ac 100644 --- a/framework/configstore/migrations_test.go +++ b/framework/configstore/migrations_test.go @@ -2547,4 +2547,119 @@ func TestMigrationMigrateProviderGovernanceToModelConfigs(t *testing.T) { Where("scope = ? AND model_name = ? AND provider = ?", tables.ModelConfigScopeGlobal, tables.ModelConfigAllModels, "openai"). Count(&count).Error) assert.Equal(t, int64(1), count, "re-run must not duplicate the wildcard config") + + // No conflict occurred, so the backup tables must not be created. + assert.False(t, db.Migrator().HasTable("governance_provider_budgets_backup"), "budget backup table should only exist on conflict") + assert.False(t, db.Migrator().HasTable("governance_provider_rate_limits_backup"), "rate limit backup table should only exist on conflict") +} + +// TestMigrationMigrateProviderGovernanceToModelConfigsConflictBackup verifies that when a +// user-created (global, provider, '*') row already occupies the wildcard slot, the provider's +// own governance rows are snapshotted into the backup tables before the FKs are cleared, and +// the pre-existing wildcard row is left untouched. +func TestMigrationMigrateProviderGovernanceToModelConfigsConflictBackup(t *testing.T) { + db := setupTestDB(t) + ctx := context.Background() + require.NoError(t, db.AutoMigrate( + &tables.TableProvider{}, &tables.TableModelConfig{}, &tables.TableBudget{}, &tables.TableRateLimit{}, + )) + + now := time.Now() + // Provider governance (would be stranded by the detach without the backup). + require.NoError(t, db.Create(&tables.TableBudget{ID: "b-provider", MaxLimit: 100, ResetDuration: "1M", LastReset: now, CreatedAt: now, UpdatedAt: now}).Error) + require.NoError(t, db.Create(&tables.TableRateLimit{ID: "rl-provider", TokenMaxLimit: schemas.Ptr(int64(1000)), TokenResetDuration: schemas.Ptr("1h"), TokenLastReset: now, RequestLastReset: now, CreatedAt: now, UpdatedAt: now}).Error) + require.NoError(t, db.Create(&tables.TableProvider{Name: "openai", BudgetID: schemas.Ptr("b-provider"), RateLimitID: schemas.Ptr("rl-provider"), CreatedAt: now, UpdatedAt: now}).Error) + // Pre-existing user-created wildcard row with its own (different) budget. + require.NoError(t, db.Create(&tables.TableBudget{ID: "b-wildcard", MaxLimit: 50, ResetDuration: "1M", LastReset: now, CreatedAt: now, UpdatedAt: now}).Error) + require.NoError(t, db.Create(&tables.TableModelConfig{ + ID: "mc-wildcard", ModelName: tables.ModelConfigAllModels, Provider: schemas.Ptr("openai"), + Scope: tables.ModelConfigScopeGlobal, BudgetID: schemas.Ptr("b-wildcard"), CreatedAt: now, UpdatedAt: now, + }).Error) + + require.NoError(t, migrationMigrateProviderGovernanceToModelConfigs(ctx, db)) + + // The pre-existing wildcard row is untouched. + var mc tables.TableModelConfig + require.NoError(t, db.Where("id = ?", "mc-wildcard").First(&mc).Error) + require.NotNil(t, mc.BudgetID) + assert.Equal(t, "b-wildcard", *mc.BudgetID) + assert.Nil(t, mc.RateLimitID) + + // Provider governance FKs are cleared. + var prov tables.TableProvider + require.NoError(t, db.Where("name = ?", "openai").First(&prov).Error) + assert.Nil(t, prov.BudgetID) + assert.Nil(t, prov.RateLimitID) + + // The stranded governance rows were snapshotted into the backup tables. + var budgetBackup struct { + ProviderName string + ID string + MaxLimit float64 + } + require.NoError(t, db.Table("governance_provider_budgets_backup"). + Select("provider_name, id, max_limit").Where("provider_name = ?", "openai"). + Scan(&budgetBackup).Error) + assert.Equal(t, "b-provider", budgetBackup.ID) + assert.Equal(t, float64(100), budgetBackup.MaxLimit) + + var rlBackup struct { + ProviderName string + ID string + } + require.NoError(t, db.Table("governance_provider_rate_limits_backup"). + Select("provider_name, id").Where("provider_name = ?", "openai"). + Scan(&rlBackup).Error) + assert.Equal(t, "rl-provider", rlBackup.ID) + + // Idempotency: re-run must not duplicate backup rows (the provider's FKs are + // already cleared, so it no longer matches the governance query). + require.NoError(t, migrationMigrateProviderGovernanceToModelConfigs(ctx, db)) + var backupCount int64 + require.NoError(t, db.Table("governance_provider_budgets_backup").Count(&backupCount).Error) + assert.Equal(t, int64(1), backupCount, "re-run must not duplicate budget backup rows") + require.NoError(t, db.Table("governance_provider_rate_limits_backup").Count(&backupCount).Error) + assert.Equal(t, int64(1), backupCount, "re-run must not duplicate rate limit backup rows") +} + +// TestMigrationMigrateProviderGovernanceToModelConfigsBackupFailureNonBlocking verifies that +// a failing governance backup is logged and skipped (rolled back to its savepoint) rather +// than failing the migration: the provider FKs are still cleared and the pre-existing +// wildcard row stays untouched. +func TestMigrationMigrateProviderGovernanceToModelConfigsBackupFailureNonBlocking(t *testing.T) { + db := setupTestDB(t) + ctx := context.Background() + require.NoError(t, db.AutoMigrate( + &tables.TableProvider{}, &tables.TableModelConfig{}, &tables.TableBudget{}, &tables.TableRateLimit{}, + )) + + now := time.Now() + require.NoError(t, db.Create(&tables.TableBudget{ID: "b-provider", MaxLimit: 100, ResetDuration: "1M", LastReset: now, CreatedAt: now, UpdatedAt: now}).Error) + require.NoError(t, db.Create(&tables.TableProvider{Name: "openai", BudgetID: schemas.Ptr("b-provider"), CreatedAt: now, UpdatedAt: now}).Error) + require.NoError(t, db.Create(&tables.TableBudget{ID: "b-wildcard", MaxLimit: 50, ResetDuration: "1M", LastReset: now, CreatedAt: now, UpdatedAt: now}).Error) + require.NoError(t, db.Create(&tables.TableModelConfig{ + ID: "mc-wildcard", ModelName: tables.ModelConfigAllModels, Provider: schemas.Ptr("openai"), + Scope: tables.ModelConfigScopeGlobal, BudgetID: schemas.Ptr("b-wildcard"), CreatedAt: now, UpdatedAt: now, + }).Error) + + // Sabotage the backup: a pre-existing table with an incompatible shape makes the + // CREATE TABLE IF NOT EXISTS a no-op and the INSERT ... SELECT fail. + require.NoError(t, db.Exec(`CREATE TABLE governance_provider_budgets_backup (bogus TEXT)`).Error) + + // The migration must still succeed. + require.NoError(t, migrationMigrateProviderGovernanceToModelConfigs(ctx, db)) + + // Provider FKs are cleared and the wildcard row is untouched despite the failed backup. + var prov tables.TableProvider + require.NoError(t, db.Where("name = ?", "openai").First(&prov).Error) + assert.Nil(t, prov.BudgetID) + var mc tables.TableModelConfig + require.NoError(t, db.Where("id = ?", "mc-wildcard").First(&mc).Error) + require.NotNil(t, mc.BudgetID) + assert.Equal(t, "b-wildcard", *mc.BudgetID) + + // The sabotaged table received no rows (the savepoint rollback discarded the attempt). + var backupCount int64 + require.NoError(t, db.Table("governance_provider_budgets_backup").Count(&backupCount).Error) + assert.Equal(t, int64(0), backupCount) }