fix(proxy): kill the whole prisma process group when a boot migration command times out (#39509)

* fix(proxy): kill the whole prisma process group when a boot migration command times out

* fix(proxy): keep boot alive when the installed proxy extras has no prisma runner

litellm/proxy/db/check_migration.py

@@ -46,17 +46,32 @@ def extract_sql_commands(diff_output: str) -> list[str]:
 
 def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
     """Checks for differences between current database and Prisma schema.
+
+    Never raises: a diff that cannot be produced, because the runner is missing,
+    because the command failed, or because it outlived its budget, is reported as
+    "no diff" so boot continues.
+
     Returns:
         A tuple containing:
         - A boolean indicating if differences were found (True) or not (False).
-        - A string with the diff output or error message.
-    Raises:
-        subprocess.CalledProcessError: If the Prisma command fails.
-        Exception: For any other errors during execution.
+        - The SQL commands that would close the diff, empty when there is none.
     """
-    verbose_logger.debug("Checking for Prisma schema diff...")
     try:
-        result: Final = subprocess.run(
+        from litellm_proxy_extras.prisma_toolchain import (
+            PRISMA_COMMAND_TIMEOUT_ENV_VAR,
+            prisma_command_timeout,
+            run_prisma,
+        )
+    except ImportError as e:
+        print(  # noqa: T201  # boot-time operator output, same channel as this helper's other messages
+            f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}"
+        )
+        return False, []
+
+    verbose_logger.debug("Checking for Prisma schema diff...")
+    timeout: Final = prisma_command_timeout()
+    try:
+        result: Final = run_prisma(
             [
                 "prisma",
                 "migrate",
@@ -67,12 +82,10 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
                 "./schema.prisma",
                 "--script",
             ],
-            capture_output=True,
-            text=True,
-            check=True,
+            timeout=timeout,
+            env=os.environ.copy(),
         )
 
-        # return True, "Migration diff generated successfully."
         sql_commands: Final = extract_sql_commands(result.stdout)
 
         if sql_commands:
@@ -83,6 +96,12 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
             return True, sql_commands
         else:
             return False, []
+    except subprocess.TimeoutExpired:
+        print(  # noqa: T201  # boot-time operator output, same channel as this helper's other messages
+            f"Timed out after {timeout}s generating the migration diff. "
+            f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer."
+        )
+        return False, []
     except subprocess.CalledProcessError as e:
         error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}"
         print(error_message)  # noqa: T201

litellm/proxy/db/prisma_client.py

@@ -937,9 +937,17 @@ class PrismaManager:
                         use_v2_resolver=use_v2_resolver,
                     )
                 else:
+                    try:
+                        from litellm_proxy_extras.prisma_toolchain import (
+                            prisma_command_timeout,
+                            run_prisma,
+                        )
+                    except ImportError as e:
+                        verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e)
+                        return False
+
                     PrismaManager._raise_if_partitioned_spend_logs()
-                    # Use prisma db push with increased timeout
-                    subprocess.run(
+                    run_prisma(
                         [
                             "prisma",
                             "db",
@@ -947,13 +955,15 @@ class PrismaManager:
                             "--accept-data-loss",
                             "--skip-generate",
                         ],
-                        timeout=60,
-                        check=True,
+                        timeout=prisma_command_timeout(),
+                        env=os.environ.copy(),
+                        stdout=None,
+                        stderr=None,
                     )
                     PrismaManager._apply_replica_identity_full_if_requested()
                     return True
-            except subprocess.TimeoutExpired:
-                verbose_proxy_logger.warning("Attempt %s timed out", attempt + 1)
+            except subprocess.TimeoutExpired as e:
+                verbose_proxy_logger.warning("Attempt %s timed out after %.0fs", attempt + 1, e.timeout)
                 time.sleep(random.randrange(5, 15))
             except subprocess.CalledProcessError as e:
                 attempts_left = 3 - attempt

tests/test_litellm/proxy/db/conftest.py

@@ -1,5 +1,11 @@
+import json
 import os
+import signal
+import sys
+import time
 from collections.abc import Generator
+from dataclasses import dataclass
+from pathlib import Path
 from typing import Optional
 
 import pytest
@@ -75,3 +81,76 @@ def reset_entra_token_provider_cache() -> Generator[None, None, None]:
 def unset_database_url(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setenv("DATABASE_URL", "about-to-be-unset")
     monkeypatch.delenv("DATABASE_URL")
+
+
+FAKE_PRISMA_CLI = """#!{python}
+import json
+import os
+import pathlib
+import subprocess
+import sys
+import time
+
+calls_file = pathlib.Path(os.environ["FAKE_PRISMA_CALLS"])
+earlier_calls = calls_file.read_text().splitlines() if calls_file.exists() else []
+with calls_file.open("a") as log:
+    print(json.dumps(sys.argv[1:]), file=log)
+if not earlier_calls and os.environ.get("FAKE_PRISMA_HANG_FIRST"):
+    grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"])
+    pathlib.Path(os.environ["FAKE_PRISMA_GRANDCHILD_PIDFILE"]).write_text(str(grandchild.pid))
+    time.sleep(600)
+sys.exit(0)
+"""
+
+
+@dataclass(frozen=True, slots=True)
+class FakePrismaCli:
+    """A stand-in `prisma` on PATH, recording every invocation.
+
+    With FAKE_PRISMA_HANG_FIRST set it hangs on its first call from a process tree
+    of its own, the way the real CLI wraps Node around a Rust schema engine, so a
+    timeout that kills only the direct child leaves the rest of that tree running.
+    """
+
+    calls_file: Path
+    grandchild_pidfile: Path
+
+    @property
+    def calls(self) -> list[list[str]]:
+        if not self.calls_file.exists():
+            return []
+        return [json.loads(line) for line in self.calls_file.read_text().splitlines()]
+
+    def grandchild_is_gone(self, within_seconds: float) -> bool:
+        deadline = time.monotonic() + within_seconds
+        while time.monotonic() < deadline:
+            try:
+                os.kill(int(self.grandchild_pidfile.read_text()), 0)
+            except ProcessLookupError:
+                return True
+            time.sleep(0.05)
+        return False
+
+
+@pytest.fixture
+def fake_prisma_cli(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[FakePrismaCli, None, None]:
+    bin_dir = tmp_path / "fakebin"
+    bin_dir.mkdir()
+    script = bin_dir / "prisma"
+    script.write_text(FAKE_PRISMA_CLI.format(python=sys.executable))
+    script.chmod(0o755)
+    cli = FakePrismaCli(
+        calls_file=tmp_path / "calls.jsonl",
+        grandchild_pidfile=tmp_path / "grandchild.pid",
+    )
+    monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}")
+    monkeypatch.setenv("FAKE_PRISMA_CALLS", str(cli.calls_file))
+    monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(cli.grandchild_pidfile))
+    monkeypatch.setenv("LITELLM_PRISMA_COMMAND_TIMEOUT", "1")
+    monkeypatch.delenv("FAKE_PRISMA_HANG_FIRST", raising=False)
+    yield cli
+    if cli.grandchild_pidfile.exists():
+        try:
+            os.kill(int(cli.grandchild_pidfile.read_text()), signal.SIGKILL)
+        except ProcessLookupError:
+            pass

tests/test_litellm/proxy/db/test_check_migration.py

@@ -37,3 +37,35 @@ def test_check_migration_out_of_sync(mocker):
     check_migration.verbose_logger.exception.assert_called_once()
     actual_message = check_migration.verbose_logger.exception.call_args[0][0]
     assert "prisma schema out of sync with db" in actual_message
+
+
+@pytest.mark.timeout(30)
+def test_migrate_diff_stops_at_its_budget_and_takes_its_process_tree_with_it(fake_prisma_cli, monkeypatch):
+    """
+    `prisma migrate diff` ran unbounded, so a database that never answers hung boot
+    before uvicorn ever started, and interrupting the proxy orphaned the schema engine.
+    """
+    from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper
+
+    monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1")
+
+    assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, [])
+    assert fake_prisma_cli.calls == [
+        ["migrate", "diff", "--from-url", "postgresql://u:p@localhost:9/x",
+         "--to-schema-datamodel", "./schema.prisma", "--script"]
+    ]
+    assert fake_prisma_cli.grandchild_is_gone(within_seconds=5)
+
+
+def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(monkeypatch):
+    """
+    Boot calls this helper directly, so an ImportError here takes the proxy down before
+    uvicorn starts. An install without the runner must lose the diagnostic, not the proxy.
+    """
+    import sys
+
+    from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper
+
+    monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None)
+
+    assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, [])

tests/test_litellm/proxy/db/test_prisma_client.py

@@ -10,7 +10,7 @@ from fastapi.testclient import TestClient
 
 
 
-from litellm.proxy.db.prisma_client import PrismaWrapper, should_update_prisma_schema
+from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, should_update_prisma_schema
 
 
 @pytest.fixture(autouse=True)
@@ -193,7 +193,10 @@ async def test_recreate_prisma_client_recovers_from_disconnected_client(
     mock_new_prisma.connect.assert_awaited_once()
 
 
-def test_db_push_applies_replica_identity_full_when_requested(monkeypatch):
+DB_PUSH_ARGV = ["db", "push", "--accept-data-loss", "--skip-generate"]
+
+
+def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_prisma_cli, unset_database_url):
     """`prisma db push` bypasses litellm-proxy-extras, so it needs its own call
     into the opt-in REPLICA IDENTITY FULL step."""
     from litellm.proxy.db.prisma_client import PrismaManager
@@ -208,14 +211,13 @@ def test_db_push_applies_replica_identity_full_when_requested(monkeypatch):
         staticmethod(lambda: applied.append(True)),
     )
 
-    with patch("litellm.proxy.db.prisma_client.subprocess.run") as mock_run:
-        assert PrismaManager.setup_database(use_migrate=False) is True
+    assert PrismaManager.setup_database(use_migrate=False) is True
 
-    assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"]
+    assert fake_prisma_cli.calls == [DB_PUSH_ARGV]
     assert applied == [True]
 
 
-def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch):
+def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch, fake_prisma_cli, unset_database_url):
     """A doc-partitioned LiteLLM_SpendLogs makes `prisma db push` rewrite the
     primary key back to ("request_id"), which Postgres rejects; the guard must
     fail fast with guidance instead of running the push."""
@@ -228,29 +230,23 @@ def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch):
     monkeypatch.setattr(
         ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True)
     )
-    with patch(  # test-quality-ok: subprocess.run is the external prisma CLI boundary, asserted never reached
-        "litellm.proxy.db.prisma_client.subprocess.run"
-    ) as mock_run:
-        with pytest.raises(RuntimeError) as err:
-            PrismaManager.setup_database(use_migrate=False)
+    with pytest.raises(RuntimeError) as err:
+        PrismaManager.setup_database(use_migrate=False)
 
     assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR
-    mock_run.assert_not_called()
+    assert fake_prisma_cli.calls == []
 
 
-def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch):
+def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch, fake_prisma_cli, unset_database_url):
     from litellm.proxy.db.prisma_client import PrismaManager
     from litellm_proxy_extras.utils import ProxyExtrasDBManager
 
     monkeypatch.setattr(
         ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False)
     )
-    with patch(  # test-quality-ok: subprocess.run is the external prisma CLI boundary, not SDK logic
-        "litellm.proxy.db.prisma_client.subprocess.run"
-    ) as mock_run:
-        assert PrismaManager.setup_database(use_migrate=False) is True
+    assert PrismaManager.setup_database(use_migrate=False) is True
 
-    assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"]
+    assert fake_prisma_cli.calls == [DB_PUSH_ARGV]
 
 
 def _entra_jwt(expires_in_seconds: int) -> str:
@@ -377,3 +373,30 @@ def test_minting_without_the_database_env_vars_names_them(azure_env, monkeypatch
 
     with pytest.raises(RuntimeError, match="DATABASE_HOST"):
         wrapper.get_rds_iam_token()
+
+
+@pytest.mark.timeout(45)
+def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_database_url, monkeypatch):
+    """
+    A timed-out `db push` used to leave Node and the schema engine writing the schema,
+    so the next attempt pushed into a database the abandoned one was still mutating.
+    """
+    monkeypatch.delenv("LITELLM_SET_REPLICA_IDENTITY_FULL", raising=False)
+    monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1")
+
+    assert PrismaManager.setup_database(use_migrate=False) is True
+    assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV]
+    assert fake_prisma_cli.grandchild_is_gone(within_seconds=5)
+
+
+def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashing_boot(
+    fake_prisma_cli, unset_database_url, monkeypatch
+):
+    """
+    An ImportError out of setup_database escapes the caller's RuntimeError handler and
+    kills boot, bypassing the operator's enforce_prisma_migration_check choice.
+    """
+    monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None)
+
+    assert PrismaManager.setup_database(use_migrate=False) is False
+    assert fake_prisma_cli.calls == []