From c14b5747377ce9c321bdd45f542b4a489f9955cc Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 14 Mar 2026 12:59:26 -0700
Subject: [PATCH] fix: prisma migrate deploy failures on pre-existing instances

Fixes failed migrations due to idempotent schema changes on pre-existing litellm instances.

Problems:
1. P3018 recovery handler never returned True on successful resolution, causing "Database setup failed after multiple retries" even when the final recovery succeeded
2. _roll_back_migration exceptions escaped the P3018 handler, preventing _resolve_specific_migration from running
3. Migration SQL used ADD COLUMN/DROP COLUMN without IF [NOT] EXISTS, failing if schema was already modified

Changes:
- Add return True after successful P3018 idempotent error recovery
- Wrap _roll_back_migration in try/except to allow recovery continuation even if rollback fails
- Make migration.sql idempotent with IF NOT EXISTS / IF EXISTS clauses

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
---
 .../20260311180521_schema_sync/migration.sql  | 14 +++++++-------
 .../20260312124619_schema_sync/migration.sql  |  2 +-
 .../litellm_proxy_extras/utils.py             | 19 +++++++++++++------
 3 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql
index 5ab834695b8..84eb70ce097 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql
@@ -1,11 +1,11 @@
 -- DropIndex
-DROP INDEX "LiteLLM_MCPServerTable_approval_status_idx";
+DROP INDEX IF EXISTS "LiteLLM_MCPServerTable_approval_status_idx";
 
 -- AlterTable
-ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN "approval_status",
-DROP COLUMN "review_notes",
-DROP COLUMN "reviewed_at",
-DROP COLUMN "source_url",
-DROP COLUMN "submitted_at",
-DROP COLUMN "submitted_by";
+ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN IF EXISTS "approval_status",
+DROP COLUMN IF EXISTS "review_notes",
+DROP COLUMN IF EXISTS "reviewed_at",
+DROP COLUMN IF EXISTS "source_url",
+DROP COLUMN IF EXISTS "submitted_at",
+DROP COLUMN IF EXISTS "submitted_by";
 
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql
index 8854fd1e205..cc48a742f20 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql
@@ -1,3 +1,3 @@
 -- AlterTable
-ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN     "models" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "models" TEXT[] DEFAULT ARRAY[]::TEXT[];
 
diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py
index f3155722187..793e2fefa8e 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/utils.py
+++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py
@@ -514,12 +514,18 @@ class ProxyExtrasDBManager:
                                 )
                                 if migration_match:
                                     migration_name = migration_match.group(1)
-                                    logger.info(
-                                        f"Rolling back migration {migration_name}"
-                                    )
-                                    ProxyExtrasDBManager._roll_back_migration(
-                                        migration_name
-                                    )
+                                    try:
+                                        logger.info(
+                                            f"Rolling back migration {migration_name}"
+                                        )
+                                        ProxyExtrasDBManager._roll_back_migration(
+                                            migration_name
+                                        )
+                                    except subprocess.CalledProcessError as rollback_err:
+                                        logger.warning(
+                                            f"Failed to roll back migration {migration_name}: {rollback_err}. "
+                                            f"It may already be in a rolled-back state."
+                                        )
                                     logger.info(
                                         f"Resolving migration {migration_name} that failed "
                                         f"due to existing schema objects"
@@ -528,6 +534,7 @@ class ProxyExtrasDBManager:
                                         migration_name
                                     )
                                     logger.info("✅ Migration resolved.")
+                                    return True
                             else:
                                 # Unknown P3018 error - log and re-raise for safety
                                 logger.warning(