From 0f45553fd6e66f29b2e846d8fe1377a46b0e1a8e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 9 Nov 2026 23:42:53 -0800 Subject: [PATCH] fix(proxy): generate the Prisma client at startup when a pip install has none A pip install ships the schema inside the package but no generated client, and since v1.83.3 no startup path generates one, so the proxy cannot start with a database. Startup now generates it from the shipped schema when it is missing, with the venv's scripts dir on PATH, and the import error names the exact command with the schema path. Fixes #26097 --- litellm/proxy/db/prisma_generate.py | 134 +++++++++++++++++ litellm/proxy/proxy_cli.py | 3 + litellm/proxy/utils.py | 7 +- tests/unit/proxy/db/test_prisma_generate.py | 153 ++++++++++++++++++++ 4 files changed, 295 insertions(+), 2 deletions(-) create mode 100644 litellm/proxy/db/prisma_generate.py create mode 100644 tests/unit/proxy/db/test_prisma_generate.py diff --git a/litellm/proxy/db/prisma_generate.py b/litellm/proxy/db/prisma_generate.py new file mode 100644 index 00000000000..d2b514a3b7d --- /dev/null +++ b/litellm/proxy/db/prisma_generate.py @@ -0,0 +1,134 @@ +"""Build the Prisma Python client from the schema LiteLLM ships, for installs that have none. + +Shipped images bake the client at build time. A pip install has only the Prisma package, and nothing in +proxy startup generates the client, so the proxy cannot import it. +""" + +import importlib +import importlib.util +import os +import subprocess +import sys +import sysconfig +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Protocol + +from litellm._logging import verbose_proxy_logger + +PRISMA_PACKAGE: Final = "prisma" +SHIPPED_SCHEMA_PATH: Final = Path(__file__).resolve().parent.parent / "schema.prisma" + + +class PrismaRunner(Protocol): + def __call__( + self, argv: Sequence[str], *, timeout: float, env: Mapping[str, str] + ) -> "subprocess.CompletedProcess[str]": ... + + +@dataclass(frozen=True, slots=True) +class ClientPresent: + pass + + +@dataclass(frozen=True, slots=True) +class ClientGenerated: + pass + + +@dataclass(frozen=True, slots=True) +class ClientGenerationFailed: + reason: str + + +PrismaClientStatus = ClientPresent | ClientGenerated | ClientGenerationFailed + + +def generate_command(schema_path: Path) -> tuple[str, ...]: + return ("prisma", "generate", "--schema", str(schema_path)) + + +def manual_generate_hint(schema_path: Path = SHIPPED_SCHEMA_PATH) -> str: + return " ".join(generate_command(schema_path)) + + +def env_with_scripts_dir(env: Mapping[str, str], scripts_dir: str) -> dict[str, str]: + """Put the interpreter's scripts dir first on PATH, so Prisma finds `prisma-client-py` in an unactivated venv.""" + path: Final = env.get("PATH") + return {**env, "PATH": os.pathsep.join((scripts_dir, path)) if path else scripts_dir} + + +def _client_is_generated(package_dir: Path) -> bool: + return (package_dir / "client.py").is_file() + + +def _refresh_loaded_package(package_name: str) -> ClientGenerationFailed | None: + """Re-run an already imported package's `__init__`, which only exposes the client when it was generated first.""" + loaded: Final = sys.modules.get(package_name) + if loaded is None: + return None + importlib.invalidate_caches() + try: + importlib.reload(loaded) + except ImportError as e: + return ClientGenerationFailed(f"the client was generated but could not be loaded: {e}") + return None + + +def ensure_prisma_client( + *, + package_dir: Path | None, + schema_path: Path, + run: PrismaRunner, + timeout: float, + env: Mapping[str, str], + package_name: str = PRISMA_PACKAGE, +) -> PrismaClientStatus: + if package_dir is None: + return ClientGenerationFailed("the prisma package is not installed") + if _client_is_generated(package_dir): + return ClientPresent() + try: + run(generate_command(schema_path), timeout=timeout, env=env) + except subprocess.TimeoutExpired: + return ClientGenerationFailed(f"prisma generate timed out after {timeout}s") + except subprocess.CalledProcessError as e: + return ClientGenerationFailed(f"prisma generate exited {e.returncode}: {e.stderr or e.stdout}") + except OSError as e: + return ClientGenerationFailed(f"could not run prisma generate: {e}") + if not _client_is_generated(package_dir): + return ClientGenerationFailed(f"prisma generate finished but wrote no client into {package_dir}") + return _refresh_loaded_package(package_name) or ClientGenerated() + + +def _installed_package_dir(package_name: str) -> Path | None: + spec: Final = importlib.util.find_spec(package_name) + if spec is None or not spec.submodule_search_locations: + return None + return Path(next(iter(spec.submodule_search_locations))) + + +def ensure_shipped_prisma_client() -> PrismaClientStatus: + """Generate the client from LiteLLM's own schema if this install has none, and log the outcome.""" + from litellm_proxy_extras.prisma_toolchain import prisma_bootstrap_timeout, run_prisma + + status: Final = ensure_prisma_client( + package_dir=_installed_package_dir(PRISMA_PACKAGE), + schema_path=SHIPPED_SCHEMA_PATH, + run=run_prisma, + timeout=prisma_bootstrap_timeout(), + env=env_with_scripts_dir(os.environ, sysconfig.get_path("scripts")), + ) + match status: + case ClientPresent(): + pass + case ClientGenerated(): + verbose_proxy_logger.info("Generated the Prisma client from %s", SHIPPED_SCHEMA_PATH) + case ClientGenerationFailed(reason=reason): + verbose_proxy_logger.warning( + "Could not generate the Prisma client (%s). Run '%s' to generate it.", + reason, + manual_generate_hint(), + ) + return status diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 41aedf4b5cd..8efef3ef6a2 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1639,6 +1639,9 @@ def run_server( PrismaManager, should_update_prisma_schema, ) + from litellm.proxy.db.prisma_generate import ensure_shipped_prisma_client + + ensure_shipped_prisma_client() if should_update_prisma_schema(general_settings.get("disable_prisma_schema_update")) is False: check_prisma_schema_diff(db_url=None) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 97c327cbe4e..cebbfe9f571 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5684,10 +5684,13 @@ class PrismaClient: from prisma import Prisma from prisma.types import DatasourceOverride except Exception as e: + from litellm.proxy.db.prisma_generate import manual_generate_hint + + generate_hint: Final = manual_generate_hint() verbose_proxy_logger.error("Failed to import Prisma client: %s", e) verbose_proxy_logger.error("This usually means 'prisma generate' hasn't been run yet.") - verbose_proxy_logger.error("Please run 'prisma generate' to generate the Prisma client.") - raise Exception("Unable to find Prisma binaries. Please run 'prisma generate' first.") + verbose_proxy_logger.error("Please run '%s' to generate the Prisma client.", generate_hint) + raise Exception(f"Unable to find Prisma binaries. Please run '{generate_hint}' first.") token_auth: Final = self.token_auth writer_token_auth: Final = None if database_url_is_pooled() else token_auth # When read-replica routing is on, tag log lines with [writer]/[reader] diff --git a/tests/unit/proxy/db/test_prisma_generate.py b/tests/unit/proxy/db/test_prisma_generate.py new file mode 100644 index 00000000000..4e6360db925 --- /dev/null +++ b/tests/unit/proxy/db/test_prisma_generate.py @@ -0,0 +1,153 @@ +import os +import subprocess +import sys +from collections.abc import Callable, Generator, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.db.prisma_generate import ( + SHIPPED_SCHEMA_PATH, + ClientGenerated, + ClientGenerationFailed, + ClientPresent, + PrismaClientStatus, + ensure_prisma_client, + env_with_scripts_dir, +) + +STUB_INIT: Final = """ +try: + from .client import Prisma +except ModuleNotFoundError: + def __getattr__(name): + raise RuntimeError("The Client hasn't been generated yet") +""" + + +@dataclass +class FakeRunner: + on_run: Callable[[], None] + calls: list[tuple[tuple[str, ...], Mapping[str, str]]] = field(default_factory=list) # mutable-ok: records calls + + def __call__( + self, argv: Sequence[str], *, timeout: float, env: Mapping[str, str] + ) -> "subprocess.CompletedProcess[str]": + self.calls.append((tuple(argv), env)) + self.on_run() + return subprocess.CompletedProcess(tuple(argv), 0, "", "") + + +def _ensure(package_dir: Path, runner: FakeRunner, package_name: str = "prisma") -> PrismaClientStatus: + return ensure_prisma_client( + package_dir=package_dir, + schema_path=SHIPPED_SCHEMA_PATH, + run=runner, + timeout=5.0, + env={"PATH": "/usr/bin"}, + package_name=package_name, + ) + + +def _write_client(package_dir: Path) -> Callable[[], None]: + return lambda: (package_dir / "client.py").write_text("class Prisma:\n pass\n") + + +def _raise(error: BaseException) -> Callable[[], None]: + def run() -> None: + raise error + + return run + + +@pytest.fixture +def stub_package(tmp_path: Path) -> Generator[tuple[str, Path], None, None]: + name: Final = "prisma_stub_for_generate_test" + package_dir: Final = tmp_path / name + package_dir.mkdir() + (package_dir / "__init__.py").write_text(STUB_INIT) + sys.path.insert(0, str(tmp_path)) + try: + yield name, package_dir + finally: + sys.path.remove(str(tmp_path)) + sys.modules.pop(name, None) + sys.modules.pop(f"{name}.client", None) + + +def test_a_missing_client_is_generated_from_the_shipped_schema(stub_package: tuple[str, Path]) -> None: + name, package_dir = stub_package + runner: Final = FakeRunner(on_run=_write_client(package_dir)) + + status: Final = _ensure(package_dir, runner, package_name=name) + + assert status == ClientGenerated() + assert [argv for argv, _ in runner.calls] == [("prisma", "generate", "--schema", str(SHIPPED_SCHEMA_PATH))] + assert SHIPPED_SCHEMA_PATH.is_file() + + +def test_an_existing_client_is_left_alone(tmp_path: Path) -> None: + _write_client(tmp_path)() + runner: Final = FakeRunner(on_run=_raise(AssertionError("must not run prisma generate"))) + + assert _ensure(tmp_path, runner) == ClientPresent() + assert runner.calls == [] + + +def test_a_package_imported_before_generation_exposes_the_new_client(stub_package: tuple[str, Path]) -> None: + name, package_dir = stub_package + stub: Final = __import__(name) + with pytest.raises(RuntimeError): + _ = stub.Prisma + + assert _ensure(package_dir, FakeRunner(on_run=_write_client(package_dir)), package_name=name) == ClientGenerated() + assert sys.modules[name].Prisma.__name__ == "Prisma" + + +@pytest.mark.parametrize( + ("error", "reason_part"), + [ + ( + subprocess.CalledProcessError(1, "prisma", "", "Could not find Prisma Schema"), + "Could not find Prisma Schema", + ), + (subprocess.TimeoutExpired("prisma", 5.0), "timed out"), + (FileNotFoundError("prisma"), "could not run"), + ], +) +def test_a_failed_generate_is_reported_not_raised(tmp_path: Path, error: BaseException, reason_part: str) -> None: + status: Final = _ensure(tmp_path, FakeRunner(on_run=_raise(error))) + + assert isinstance(status, ClientGenerationFailed) + assert reason_part in status.reason + + +def test_a_generate_that_writes_no_client_is_a_failure(tmp_path: Path) -> None: + status: Final = _ensure(tmp_path, FakeRunner(on_run=lambda: None)) + + assert isinstance(status, ClientGenerationFailed) + + +def test_a_missing_prisma_package_is_a_failure() -> None: + runner: Final = FakeRunner(on_run=_raise(AssertionError("must not run prisma generate"))) + status: Final = ensure_prisma_client( + package_dir=None, schema_path=SHIPPED_SCHEMA_PATH, run=runner, timeout=5.0, env={} + ) + + assert isinstance(status, ClientGenerationFailed) + assert runner.calls == [] + + +@pytest.mark.parametrize( + ("env", "expected_path"), + [ + ({"PATH": "/usr/bin", "HOME": "/root"}, os.pathsep.join(("/opt/litellm/.venv/bin", "/usr/bin"))), + ({"HOME": "/root"}, "/opt/litellm/.venv/bin"), + ], +) +def test_the_venv_scripts_dir_comes_first_on_path(env: dict[str, str], expected_path: str) -> None: + result: Final = env_with_scripts_dir(env, "/opt/litellm/.venv/bin") + + assert result == {**env, "PATH": expected_path}