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
This commit is contained in:
ryan-crabbe-berri 2026-11-09 23:42:53 -08:00
parent b3f75884a6
commit 0f45553fd6
4 changed files with 295 additions and 2 deletions

View file

@ -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

View file

@ -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)

View file

@ -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]

View file

@ -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}