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}