From 9049f3786448e16acbccc8cb5c560c11883c50ea Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Tue, 21 Apr 2026 15:34:24 -0700
Subject: [PATCH] [Fix] v2 migration resolver: address Greptile review findings
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Open the psycopg connection in `_warn_if_db_ahead_of_head` with
  autocommit=True. Without it, psycopg3's `with conn` calls COMMIT on
  clean exit, which fails after the `UndefinedTable` (fresh-DB) branch
  left the transaction in an aborted state — crashing first-run startups.

- Wrap the v2 `prisma db push` path in try/except and raise RuntimeError
  on CalledProcessError/TimeoutExpired. Otherwise these propagate past
  proxy_cli.py's `except RuntimeError` as unhandled tracebacks.

- Reword the loop-exhaustion error to cover the non-timeout exit path
  (repeated P3005/P3009/P3018 idempotent-recovery `continue`s), not
  just persistent timeouts.

Adds a unit test for the db_push error wrapping.
---
 .../litellm_proxy_extras/utils.py             | 21 ++++++++++++++++---
 .../tests/test_setup_database_fail_fast.py    | 12 +++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py
index 04005ce2548..a234ad18ba8 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/utils.py
+++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py
@@ -456,7 +456,13 @@ class ProxyExtrasDBManager:
         known = set(ProxyExtrasDBManager._get_migration_names(migrations_dir))
 
         try:
-            with psycopg.connect(cleaned_url, connect_timeout=10) as conn:
+            # autocommit=True keeps the SELECT outside a transaction. Without
+            # it, psycopg3's `with conn` calls COMMIT on clean exit — which
+            # fails after `UndefinedTable` (fresh DB) leaves the transaction
+            # in an aborted state.
+            with psycopg.connect(
+                cleaned_url, connect_timeout=10, autocommit=True
+            ) as conn:
                 try:
                     rows = conn.execute(
                         "SELECT migration_name FROM _prisma_migrations "
@@ -521,6 +527,13 @@ class ProxyExtrasDBManager:
                     env=_get_prisma_env(),
                 )
                 return True
+            except (
+                subprocess.CalledProcessError,
+                subprocess.TimeoutExpired,
+            ) as e:
+                # Re-raise as RuntimeError so proxy_cli.py's
+                # `except RuntimeError` catches it and exits cleanly.
+                raise RuntimeError(f"prisma db push failed.\n\nDetail: {e}") from e
             finally:
                 os.chdir(original_dir)
 
@@ -624,8 +637,10 @@ class ProxyExtrasDBManager:
                     ) from e
 
             raise RuntimeError(
-                "Database migration failed after 4 attempts (persistent timeouts). "
-                "Check database connectivity and load."
+                "Database migration failed after 4 attempts (retry loop "
+                "exhausted by timeouts or repeated idempotent-recovery "
+                "continues). Check database connectivity, load, and "
+                "_prisma_migrations ledger state."
             )
         finally:
             os.chdir(original_dir)
diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py
index 38df390a42d..573137c90ed 100644
--- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py
+++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py
@@ -132,6 +132,18 @@ def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
     assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path"
 
 
+def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
+    """v2: a failing `prisma db push` must raise RuntimeError, not leak
+    CalledProcessError past proxy_cli.py's `except RuntimeError`."""
+    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
+    (tmp_path / "schema.prisma").write_text("// stub")
+
+    stderr = "db push error"
+    with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
+        with pytest.raises(RuntimeError, match="prisma db push failed"):
+            ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
+
+
 def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
     """v2 must never call _resolve_all_migrations — that's the bug it fixes."""
     monkeypatch.setattr(