fix(proxy-extras): stop a db push timeout crashing the migration job

subprocess.run leaves stderr as bytes on TimeoutExpired even under
text=True, unlike CalledProcessError. Classifying both in one handler
meant a real `prisma db push` timeout died on a TypeError, which
proxy_cli.py's `except RuntimeError` does not catch, so the migrations
Job container ended on an unhandled traceback instead of a clean exit.

Give the timeout its own handler and retry it, matching what the migrate
deploy loop beside it already does. That puts a fallthrough back into the
loop, so the trailing raise removed in the previous commit is reachable
again and comes back with it.

Also drop a comment restating why the resolver cases exist and widen the
db push test's docstring, which had stopped describing what it covers.

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -709,10 +709,13 @@ class ProxyExtrasDBManager:
                             env=_get_prisma_env(),
                         )
                         return True
-                    except (
-                        subprocess.CalledProcessError,
-                        subprocess.TimeoutExpired,
-                    ) as e:
+                    except subprocess.TimeoutExpired:
+                        logger.info(
+                            "prisma db push attempt %s timed out, retrying",
+                            attempt + 1,
+                        )
+                        time.sleep(random.randrange(5, 15))
+                    except subprocess.CalledProcessError as e:
                         stderr = e.stderr or ""
                         transient = ProxyExtrasDBManager._transient_prisma_failure(
                             stderr
@@ -732,6 +735,9 @@ class ProxyExtrasDBManager:
                             stderr,
                         )
                         time.sleep(random.randrange(5, 15))
+                raise RuntimeError(
+                    f"prisma db push failed after {_PRISMA_ATTEMPTS} attempts."
+                )
             finally:
                 os.chdir(original_dir)
 

tests/litellm-proxy-extras/test_setup_database_fail_fast.py

@@ -460,6 +460,92 @@ def test_v2_db_push_retries_are_bounded_and_report_the_prisma_error(
     assert "P1001" in str(exc.value)
 
 
+def _db_push_only(push_side_effect):
+    """subprocess.run stand-in that only intercepts `prisma db push`."""
+    pushes = {"n": 0}
+
+    def _run(*args, **kwargs):
+        cmd = list(args[0] if args else kwargs.get("args", []))
+        if cmd[-3:] != ["db", "push", "--accept-data-loss"]:
+            return _DeployApplied()
+        pushes["n"] += 1
+        return push_side_effect(pushes["n"], cmd)
+
+    return _run, pushes
+
+
+def _timed_out_for_real():
+    """Capture what subprocess.run really puts on a TimeoutExpired.
+
+    Under text=True it still leaves stderr as bytes, unlike CalledProcessError,
+    so hardcoding a str here would test a shape production never sees. Derived
+    at import, before any test patches subprocess.run.
+    """
+    try:
+        subprocess.run(
+            ["sh", "-c", "echo 'Error: P1001 unreachable' >&2; sleep 5"],
+            timeout=0.2,
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except subprocess.TimeoutExpired as e:
+        return e
+    raise AssertionError("the helper command was supposed to time out")
+
+
+_TIMEOUT_TEMPLATE = _timed_out_for_real()
+
+
+def _real_timeout_expired(cmd):
+    return subprocess.TimeoutExpired(
+        cmd=cmd,
+        timeout=_TIMEOUT_TEMPLATE.timeout,
+        output=_TIMEOUT_TEMPLATE.stdout,
+        stderr=_TIMEOUT_TEMPLATE.stderr,
+    )
+
+
+def test_v2_db_push_retries_a_timeout(monkeypatch, tmp_path):
+    """v2: a `prisma db push` that times out is retried, not turned into a
+    TypeError by classifying its bytes stderr as if it were text."""
+    _prepare_v2_resolver(monkeypatch, tmp_path)
+
+    def _side_effect(n, cmd):
+        if n == 1:
+            raise _real_timeout_expired(cmd)
+        return _DeployApplied()
+
+    monkeypatch.setattr(
+        ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False
+    )
+    run, pushes = _db_push_only(_side_effect)
+    with patch("subprocess.run", side_effect=run):
+        ok = ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
+
+    assert ok is True
+    assert pushes["n"] == 2
+
+
+def test_v2_db_push_timeouts_are_bounded(monkeypatch, tmp_path):
+    """v2: a `prisma db push` that never stops timing out gives up as a
+    RuntimeError, which is the only exception proxy_cli.py exits cleanly on."""
+    _prepare_v2_resolver(monkeypatch, tmp_path)
+
+    def _side_effect(n, cmd):
+        raise _real_timeout_expired(cmd)
+
+    monkeypatch.setattr(
+        ProxyExtrasDBManager, "spend_logs_is_partitioned", lambda: False
+    )
+    run, pushes = _db_push_only(_side_effect)
+    with patch("subprocess.run", side_effect=run):
+        with pytest.raises(RuntimeError, match=r"prisma db push failed after \d+"):
+            ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
+
+    assert pushes["n"] == _PRISMA_ATTEMPTS
+
+
 def test_v2_unclassified_failure_is_not_treated_as_transient(monkeypatch, tmp_path):
     """v2: an unrecognised deploy failure still raises on the first attempt."""
     _prepare_v2_resolver(monkeypatch, tmp_path)

tests/test_litellm/proxy/test_proxy_cli.py

@@ -1737,7 +1737,8 @@ class TestRunServerDbSetup:
         mock_atexit_register,
         mock_subprocess_run,
     ):
-        """Test that use_prisma_db_push flag correctly controls PrismaManager.setup_database use_migrate parameter"""
+        """Which resolver and which migration mode run_server hands setup_database,
+        across the db push flag, the v2/legacy flag pair and USE_V2_MIGRATION_RESOLVER."""
         from litellm.proxy.proxy_cli import run_server
 
         # Mock subprocess.run to simulate prisma being available
@@ -1805,10 +1806,6 @@ class TestRunServerDbSetup:
                 use_migrate=False, use_v2_resolver=True
             )
 
-            # Test 3+: the resolver default and both routes back to v1. The flag
-            # covers a CLI boot; USE_V2_MIGRATION_RESOLVER covers deploys that
-            # cannot pass one, where prisma_migration.py fixes the argv. An
-            # explicit flag beats the env var.
             for argv, env_value, expected_v2 in (
                 ([], None, True),
                 (["--use_v2_migration_resolver"], None, True),