fix: supports backup for governance migration

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=<name>, 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)
 						}
 					}
 				}

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)
 }