From 221a2b415262c63dbf4ba28271a621cb8bfa99c4 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 10 Nov 2026 00:22:06 -0800 Subject: [PATCH] fix(proxy): keep the Prisma generate helper inside the lint gates Stream prisma's own output to the proxy log, type the extras toolchain through a small protocol, and avoid sys.path edits in the test. --- litellm/proxy/db/prisma_generate.py | 30 +++++++++------- tests/unit/proxy/db/test_prisma_generate.py | 38 ++++++++++++--------- 2 files changed, 40 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/db/prisma_generate.py b/litellm/proxy/db/prisma_generate.py index d2b514a3b7d..7f23f3c262e 100644 --- a/litellm/proxy/db/prisma_generate.py +++ b/litellm/proxy/db/prisma_generate.py @@ -13,7 +13,8 @@ import sysconfig from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path -from typing import Final, Protocol +from types import MappingProxyType +from typing import Final, Protocol, TypeAlias from litellm._logging import verbose_proxy_logger @@ -23,7 +24,13 @@ SHIPPED_SCHEMA_PATH: Final = Path(__file__).resolve().parent.parent / "schema.pr class PrismaRunner(Protocol): def __call__( - self, argv: Sequence[str], *, timeout: float, env: Mapping[str, str] + self, + argv: Sequence[str], + *, + timeout: float, + env: Mapping[str, str], + stdout: int | None, + stderr: int | None, ) -> "subprocess.CompletedProcess[str]": ... @@ -42,7 +49,7 @@ class ClientGenerationFailed: reason: str -PrismaClientStatus = ClientPresent | ClientGenerated | ClientGenerationFailed +PrismaClientStatus: TypeAlias = ClientPresent | ClientGenerated | ClientGenerationFailed def generate_command(schema_path: Path) -> tuple[str, ...]: @@ -53,10 +60,10 @@ 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]: +def env_with_scripts_dir(env: Mapping[str, str], scripts_dir: str) -> Mapping[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} + return MappingProxyType({**env, "PATH": os.pathsep.join((scripts_dir, path)) if path else scripts_dir}) def _client_is_generated(package_dir: Path) -> bool: @@ -90,11 +97,11 @@ def ensure_prisma_client( if _client_is_generated(package_dir): return ClientPresent() try: - run(generate_command(schema_path), timeout=timeout, env=env) + run(generate_command(schema_path), timeout=timeout, env=env, stdout=None, stderr=None) 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}") + return ClientGenerationFailed(f"prisma generate exited {e.returncode}, see its output above") except OSError as e: return ClientGenerationFailed(f"could not run prisma generate: {e}") if not _client_is_generated(package_dir): @@ -109,15 +116,15 @@ def _installed_package_dir(package_name: str) -> Path | None: return Path(next(iter(spec.submodule_search_locations))) -def ensure_shipped_prisma_client() -> PrismaClientStatus: +def ensure_shipped_prisma_client() -> None: """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 + from litellm_proxy_extras import prisma_toolchain # pyright: ignore[reportUnknownVariableType] # untyped in CI env status: Final = ensure_prisma_client( package_dir=_installed_package_dir(PRISMA_PACKAGE), schema_path=SHIPPED_SCHEMA_PATH, - run=run_prisma, - timeout=prisma_bootstrap_timeout(), + run=prisma_toolchain.run_prisma, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped in CI env + timeout=prisma_toolchain.prisma_bootstrap_timeout(), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped in CI env env=env_with_scripts_dir(os.environ, sysconfig.get_path("scripts")), ) match status: @@ -131,4 +138,3 @@ def ensure_shipped_prisma_client() -> PrismaClientStatus: reason, manual_generate_hint(), ) - return status diff --git a/tests/unit/proxy/db/test_prisma_generate.py b/tests/unit/proxy/db/test_prisma_generate.py index 4e6360db925..d420b102335 100644 --- a/tests/unit/proxy/db/test_prisma_generate.py +++ b/tests/unit/proxy/db/test_prisma_generate.py @@ -1,3 +1,4 @@ +import importlib import os import subprocess import sys @@ -33,7 +34,13 @@ class FakeRunner: 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] + self, + argv: Sequence[str], + *, + timeout: float, + env: Mapping[str, str], + stdout: int | None, + stderr: int | None, ) -> "subprocess.CompletedProcess[str]": self.calls.append((tuple(argv), env)) self.on_run() @@ -52,7 +59,10 @@ def _ensure(package_dir: Path, runner: FakeRunner, package_name: str = "prisma") def _write_client(package_dir: Path) -> Callable[[], None]: - return lambda: (package_dir / "client.py").write_text("class Prisma:\n pass\n") + def run() -> None: + _ = (package_dir / "client.py").write_text("class Prisma:\n pass\n") + + return run def _raise(error: BaseException) -> Callable[[], None]: @@ -63,18 +73,17 @@ def _raise(error: BaseException) -> Callable[[], None]: @pytest.fixture -def stub_package(tmp_path: Path) -> Generator[tuple[str, Path], None, None]: +def stub_package(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> 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)) + _ = (package_dir / "__init__.py").write_text(STUB_INIT) + monkeypatch.setattr(sys, "path", [str(tmp_path), *sys.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) + _ = 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: @@ -98,21 +107,18 @@ def test_an_existing_client_is_left_alone(tmp_path: Path) -> None: 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) + stub: Final = importlib.import_module(name) with pytest.raises(RuntimeError): - _ = stub.Prisma + _ = hasattr(stub, "Prisma") assert _ensure(package_dir, FakeRunner(on_run=_write_client(package_dir)), package_name=name) == ClientGenerated() - assert sys.modules[name].Prisma.__name__ == "Prisma" + assert "Prisma" in vars(sys.modules[name]) @pytest.mark.parametrize( ("error", "reason_part"), [ - ( - subprocess.CalledProcessError(1, "prisma", "", "Could not find Prisma Schema"), - "Could not find Prisma Schema", - ), + (subprocess.CalledProcessError(1, "prisma"), "exited 1"), (subprocess.TimeoutExpired("prisma", 5.0), "timed out"), (FileNotFoundError("prisma"), "could not run"), ], @@ -150,4 +156,4 @@ def test_a_missing_prisma_package_is_a_failure() -> None: 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} + assert dict(result) == {**env, "PATH": expected_path}