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

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

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}