From b44099c1e8752e4cb12d16039ab8c1b31514d63f Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 14 Mar 2026 15:50:08 -0700 Subject: [PATCH] fix: retry after resolving idempotent migration, guard DROP CONSTRAINT Three fixes: 1. Both P3009 and P3018 idempotent handlers returned True after resolving a single migration, exiting before remaining pending migrations were applied. Now they continue the retry loop so prisma migrate deploy runs again for any remaining migrations. 2. Two migration files had bare DROP CONSTRAINT without a DO $$ IF EXISTS guard, which fails if the constraint was already dropped. Wrapped both in idempotent DO $$ blocks. 3. Added test_drop_constraint_is_guarded to catch unguarded DROP CONSTRAINT in future migrations. Co-Authored-By: Claude Opus 4.6 --- .../migration.sql | 7 ++++++- .../migration.sql | 7 ++++++- .../litellm_proxy_extras/utils.py | 9 +++++---- .../test_litellm_proxy_extras_utils.py | 18 ++++++++++++++++++ 4 files changed, 35 insertions(+), 6 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql index 0fbb5024dde..ef53fbf3f4a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql @@ -1,5 +1,10 @@ -- DropForeignKey -ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey"; +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey"; + END IF; +END $$; -- AddForeignKey DO $$ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql index 3c113a21f79..75d7e0e74b1 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql @@ -1,5 +1,10 @@ -- DropForeignKey -ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey"; +DO $$ +BEGIN + IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey"; + END IF; +END $$; -- AlterTable ALTER TABLE "LiteLLM_ManagedFileTable" ALTER COLUMN "file_object" DROP NOT NULL; diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 7a14a109530..5b4feed2564 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -422,9 +422,8 @@ class ProxyExtrasDBManager: failed_migration ) logger.info( - f"✅ Migration {failed_migration} resolved." + f"✅ Migration {failed_migration} resolved, retrying to apply remaining migrations" ) - return True else: logger.info( f"Found failed migration: {failed_migration}, marking as rolled back" @@ -534,8 +533,10 @@ class ProxyExtrasDBManager: ProxyExtrasDBManager._resolve_specific_migration( migration_name ) - logger.info("✅ Migration resolved.") - return True + logger.info( + f"✅ Migration {migration_name} resolved, " + f"retrying to apply remaining migrations" + ) except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as resolve_err: logger.warning( f"Failed to resolve migration {migration_name}: {resolve_err}" diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 1fc93cd4966..05f4b8a98e1 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -275,3 +275,21 @@ class TestMigrationSQLIdempotency: "ADD CONSTRAINT without DO $$ IF NOT EXISTS guard found in migrations:\n" + "\n".join(violations) ) + + def test_drop_constraint_is_guarded(self, all_migrations): + """DROP CONSTRAINT must be inside a DO $$ IF EXISTS block""" + violations = [] + for migration_name, sql in all_migrations: + lines = sql.splitlines() + in_do_block = False + for line_num, line in enumerate(lines, 1): + if re.search(r"DO\s+\$\$", line, re.IGNORECASE): + in_do_block = True + if re.search(r"END\s+\$\$", line, re.IGNORECASE): + in_do_block = False + if re.search(r"DROP\s+CONSTRAINT\s+", line, re.IGNORECASE) and not in_do_block: + violations.append(f" {migration_name}:{line_num}: {line.strip()}") + assert not violations, ( + "DROP CONSTRAINT without DO $$ IF EXISTS guard found in migrations:\n" + + "\n".join(violations) + )