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