mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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:
parent
b3f75884a6
commit
0f45553fd6
4 changed files with 295 additions and 2 deletions
134
litellm/proxy/db/prisma_generate.py
Normal file
134
litellm/proxy/db/prisma_generate.py
Normal 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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
153
tests/unit/proxy/db/test_prisma_generate.py
Normal file
153
tests/unit/proxy/db/test_prisma_generate.py
Normal 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}
|
||||
Loading…
Add table
Reference in a new issue