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 <noreply@anthropic.com>

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 $$

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;

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

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