From 46f88198ce9a549e6a18e80b72ed96a287999c74 Mon Sep 17 00:00:00 2001
From: jesus <jesus@berri.ai>
Date: Wed, 28 Oct 2026 00:59:18 +0000
Subject: [PATCH] fix(proxy): skip the migration replay when the ledger already
 records the files

A failed migrate diff during the post-deploy sanity check (mark_all_applied=False)
replayed every migration file onto an up-to-date database, which recreated the
renamed LiteLLM_Engine tables and the dropped (prompt_id, version) unique index.
The replay now only runs on the legacy no-ledger path, and a missing relation in
the replay counts as a real failure instead of an idempotent one
---
 .../litellm_proxy_extras/utils.py             | 10 +++++++
 .../test_litellm_proxy_extras_utils.py        | 27 ++++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py
index 3f825d2b026..aacbed6b8e1 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/utils.py
+++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py
@@ -146,6 +146,7 @@ class _MigrateAttemptBudget:
 
 _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
 _SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE)
+_MISSING_RELATION: Final = re.compile(r'(?<!of )relation "[^"]+" does not exist')
 _SPEND_LOGS_PK_CLAUSE_RE = re.compile(
     r'^(?:DROP\s+CONSTRAINT\s+"[^"]*_pkey"'
     r'|ADD\s+(?:CONSTRAINT\s+"[^"]*"\s+)?PRIMARY\s+KEY\s*\([^)]*\))$',
@@ -550,6 +551,13 @@ class ProxyExtrasDBManager:
 
         # check if the migration was created
         if not diff_sql_path.exists():
+            if not mark_all_applied:
+                logger.warning(
+                    "Migration diff was not created (prisma migrate diff failed, likely a pooler URL). "
+                    "The ledger already records every migration file, so skipping the drift check "
+                    "instead of replaying them"
+                )
+                return
             logger.warning(
                 "Migration diff was not created (prisma migrate diff failed — "
                 "likely a pooler URL). Falling back to direct SQL execution of "
@@ -633,6 +641,8 @@ class ProxyExtrasDBManager:
         except subprocess.CalledProcessError as e:
             stderr: Final = str(e.stderr or "")  # pyright: ignore[reportAny]  # CalledProcessError.stderr is Any
             logger.warning(f"Failed to apply migration {mig_file.parent.name}: {stderr}")
+            if _MISSING_RELATION.search(stderr):
+                return False
             return ProxyExtrasDBManager._is_idempotent_error(stderr)
         except subprocess.TimeoutExpired:
             logger.warning(f"Migration {mig_file.parent.name} timed out.")
diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py
index 4e6ec4e9084..f3d215f81db 100644
--- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py
+++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py
@@ -670,6 +670,7 @@ class TestResolveAllMigrationsLedger:
         execute_fails,
         diff_fails: bool = False,
         execute_stderr: str = "boom",
+        mark_all_applied: bool = True,
     ):
         import subprocess as subprocess_module
 
@@ -711,7 +712,7 @@ class TestResolveAllMigrationsLedger:
             return _FakeCompleted()
 
         monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", fake_run)
-        ProxyExtrasDBManager._resolve_all_migrations(str(tmp_path), "schema.prisma")
+        ProxyExtrasDBManager._resolve_all_migrations(str(tmp_path), "schema.prisma", mark_all_applied=mark_all_applied)
         return calls
 
     def _resolved(self, calls):
@@ -764,6 +765,30 @@ class TestResolveAllMigrationsLedger:
         )
         assert self._resolved(calls) == []
 
+    def test_failed_diff_with_a_missing_table_marks_nothing_applied(self, monkeypatch, tmp_path):
+        calls = self._run(
+            monkeypatch, tmp_path, partitioned=False, execute_fails=True, diff_fails=True,
+            execute_stderr='ERROR: relation "LiteLLM_TeamTable" does not exist',
+        )
+        assert self._resolved(calls) == []
+
+    def test_failed_diff_with_an_already_dropped_constraint_still_marks_migrations_applied(self, monkeypatch, tmp_path):
+        calls = self._run(
+            monkeypatch, tmp_path, partitioned=False, execute_fails=True, diff_fails=True,
+            execute_stderr=(
+                'ERROR: constraint "LiteLLM_TeamMembership_budget_id_fkey" '
+                'of relation "LiteLLM_TeamMembership" does not exist'
+            ),
+        )
+        assert len(self._resolved(calls)) == 1
+
+    def test_failed_diff_during_the_post_deploy_sanity_check_replays_nothing(self, monkeypatch, tmp_path):
+        calls = self._run(
+            monkeypatch, tmp_path, partitioned=False, execute_fails=False, diff_fails=True, mark_all_applied=False
+        )
+        assert [c for c in calls if isinstance(c, list) and "execute" in c] == []
+        assert self._resolved(calls) == []
+
 
 class TestPartitionedSpendLogsPushGuard:
     def _forbid_subprocess(self, monkeypatch):