diff --git a/internal/database/postgres/migrations/000040_split_savingsplans.down.sql b/internal/database/postgres/migrations/000040_split_savingsplans.down.sql index 41149a2e8..b4742ff22 100644 --- a/internal/database/postgres/migrations/000040_split_savingsplans.down.sql +++ b/internal/database/postgres/migrations/000040_split_savingsplans.down.sql @@ -65,20 +65,24 @@ SET services = ( -- created post-split with only sagemaker/database/ec2instance -- still gets a usable umbrella key on rollback. SELECT 'aws:savings-plans' AS new_key, v AS new_val - FROM jsonb_each(services) AS e(k, v) - WHERE k IN ( - 'aws:savings-plans-compute', - 'aws:savings-plans-ec2instance', - 'aws:savings-plans-sagemaker', - 'aws:savings-plans-database' - ) - ORDER BY CASE k - WHEN 'aws:savings-plans-compute' THEN 1 - WHEN 'aws:savings-plans-ec2instance' THEN 2 - WHEN 'aws:savings-plans-sagemaker' THEN 3 - ELSE 4 - END - LIMIT 1 + FROM ( + SELECT v, + CASE k + WHEN 'aws:savings-plans-compute' THEN 1 + WHEN 'aws:savings-plans-ec2instance' THEN 2 + WHEN 'aws:savings-plans-sagemaker' THEN 3 + ELSE 4 + END AS priority + FROM jsonb_each(services) AS e(k, v) + WHERE k IN ( + 'aws:savings-plans-compute', + 'aws:savings-plans-ec2instance', + 'aws:savings-plans-sagemaker', + 'aws:savings-plans-database' + ) + ORDER BY priority + LIMIT 1 + ) sp_best ) merged ) WHERE services ?| ARRAY[ diff --git a/internal/database/postgres/migrations/helpers_test.go b/internal/database/postgres/migrations/helpers_test.go new file mode 100644 index 000000000..9a76f93b0 --- /dev/null +++ b/internal/database/postgres/migrations/helpers_test.go @@ -0,0 +1,17 @@ +//go:build integration +// +build integration + +package migrations_test + +import ( + "path/filepath" + "runtime" +) + +// getMigrationsPath resolves the migrations directory relative to this test +// file so that migration test files can locate the SQL files regardless of +// the working directory when tests are invoked. +func getMigrationsPath() string { + _, filename, _, _ := runtime.Caller(0) + return filepath.Dir(filename) +} diff --git a/internal/database/postgres/migrations/savings_snapshots_pk_test.go b/internal/database/postgres/migrations/savings_snapshots_pk_test.go new file mode 100644 index 000000000..128c86332 --- /dev/null +++ b/internal/database/postgres/migrations/savings_snapshots_pk_test.go @@ -0,0 +1,111 @@ +//go:build integration +// +build integration + +package migrations_test + +import ( + "context" + "testing" + + "github.com/LeanerCloud/CUDly/internal/database/postgres/migrations" + "github.com/LeanerCloud/CUDly/internal/database/postgres/testhelpers" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestMigration_SavingsSnapshotsPK verifies that migration 000027 is +// idempotent on a fresh database. +// +// On a fresh DB, migration 000018 already adds savings_snapshots_pkey. Before +// the fix, 000027's bare ADD CONSTRAINT would fail with "multiple primary +// keys for table savings_snapshots". After the fix, the DROP CONSTRAINT IF +// EXISTS guard makes 000027 safe to apply regardless of whether 000018 has +// already created the constraint. +func TestMigration_SavingsSnapshotsPK(t *testing.T) { + ctx := context.Background() + migrationsPath := getMigrationsPath() + + t.Run("full migration from scratch succeeds (000018 + 000027 coexist)", func(t *testing.T) { + container, err := testhelpers.SetupPostgresContainer(ctx, t) + require.NoError(t, err) + defer container.Cleanup(ctx) + pool := container.DB.Pool() + + // Running all migrations must succeed on a fresh DB — 000027's + // DROP CONSTRAINT IF EXISTS prevents "multiple primary keys" + // that the pre-fix ADD CONSTRAINT caused. + require.NoError(t, migrations.RunMigrations(ctx, pool, migrationsPath, "", ""), + "full migration from scratch should succeed with idempotent 000027") + + // The primary key must exist and cover (id, timestamp). + var constraintExists bool + err = pool.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 FROM pg_constraint + WHERE conname = 'savings_snapshots_pkey' + AND conrelid = 'savings_snapshots'::regclass + AND contype = 'p' + )`).Scan(&constraintExists) + require.NoError(t, err) + assert.True(t, constraintExists, "savings_snapshots_pkey should exist after migrations") + }) + + t.Run("rollback to 000026 then re-apply 000027 succeeds", func(t *testing.T) { + container, err := testhelpers.SetupPostgresContainer(ctx, t) + require.NoError(t, err) + defer container.Cleanup(ctx) + pool := container.DB.Pool() + + require.NoError(t, migrations.RunMigrations(ctx, pool, migrationsPath, "", "")) + + // Count how many applied migrations are above version 26. + // Using COUNT from schema_migrations is correct even when migration + // numbering has gaps: Steps(-n) steps back through n applied + // migrations, so we need the actual row count, not an arithmetic + // difference between version numbers. + var stepsToRollback int + err = pool.QueryRow(ctx, `SELECT COUNT(*) FROM schema_migrations WHERE version > 26`).Scan(&stepsToRollback) + require.NoError(t, err) + require.Greater(t, stepsToRollback, 0, "there should be migrations above 000026") + + // RollbackMigrations caps a single call at 10 steps; call in a loop. + const maxRollbackPerCall = 10 + for remaining := stepsToRollback; remaining > 0; remaining -= maxRollbackPerCall { + batch := remaining + if batch > maxRollbackPerCall { + batch = maxRollbackPerCall + } + require.NoError(t, migrations.RollbackMigrations(ctx, pool, migrationsPath, batch), + "rollback to 000026 should succeed") + } + + // Verify the constraint is gone (000027 was rolled back). + var constraintExists bool + err = pool.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 FROM pg_constraint + WHERE conname = 'savings_snapshots_pkey' + AND conrelid = 'savings_snapshots'::regclass + AND contype = 'p' + )`).Scan(&constraintExists) + require.NoError(t, err) + assert.False(t, constraintExists, "savings_snapshots_pkey should be gone after rollback past 000027") + + // Re-apply all migrations — 000027 must succeed even though + // 000018's schema_migrations entry is still present (that migration + // was not rolled back, though its original constraint was superseded + // when 000027 ran and dropped/re-created savings_snapshots_pkey). + require.NoError(t, migrations.RunMigrations(ctx, pool, migrationsPath, "", ""), + "re-applying 000027 over a DB where 000018 already ran must succeed") + + err = pool.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 FROM pg_constraint + WHERE conname = 'savings_snapshots_pkey' + AND conrelid = 'savings_snapshots'::regclass + AND contype = 'p' + )`).Scan(&constraintExists) + require.NoError(t, err) + assert.True(t, constraintExists, "savings_snapshots_pkey should be restored after re-apply") + }) +} diff --git a/internal/database/postgres/migrations/split_savingsplans_test.go b/internal/database/postgres/migrations/split_savingsplans_test.go index 2c6843624..81fb905bb 100644 --- a/internal/database/postgres/migrations/split_savingsplans_test.go +++ b/internal/database/postgres/migrations/split_savingsplans_test.go @@ -5,8 +5,6 @@ package migrations_test import ( "context" - "path/filepath" - "runtime" "testing" "github.com/LeanerCloud/CUDly/internal/database/postgres/migrations" @@ -16,12 +14,6 @@ import ( "github.com/stretchr/testify/require" ) -// getMigrationsPath resolves the migrations directory next to this test file. -func getMigrationsPath() string { - _, filename, _, _ := runtime.Caller(0) - return filepath.Dir(filename) -} - // TestMigration_SplitSavingsPlans drives the 000040 migration through // the three scenarios from the plan §7: umbrella-only, umbrella + PR // #71's sagemaker row, and fresh install. Round-trips up/down to verify