Merge pull request #40768 from BerriAI/litellm_prisma_cli_resolve

fix(proxy): run migrations through python -m prisma when the prisma console script is not on PATH

litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py

@@ -37,11 +37,13 @@ raised it above the deploy default keeps that larger budget for deploy unless
 the deploy override says otherwise.
 """
 
+import importlib.util
 import math
 import os
 import shutil
 import signal
 import subprocess
+import sys
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from pathlib import Path
@@ -64,6 +66,7 @@ DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
 DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
 
 BOOTSTRAP_ARG = "--version"
+PRISMA_CONSOLE_SCRIPT = "prisma"
 
 
 @dataclass(frozen=True)
@@ -184,6 +187,28 @@ def _kill_process_group(process: "subprocess.Popen[str]") -> None:
         return
 
 
+def prisma_cli_available() -> bool:
+    """Whether some way of running the Prisma CLI exists: the console script on PATH or the importable package."""
+    if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None:
+        return True
+    return importlib.util.find_spec(PRISMA_CONSOLE_SCRIPT) is not None
+
+
+def resolve_prisma_argv(argv: Sequence[str]) -> tuple[str, ...]:
+    """Route a bare ``prisma`` command through ``python -m prisma`` when the console script is not on PATH.
+
+    The console script and ``python -m prisma`` are the same entry point, but
+    only the module form survives an interpreter whose ``bin`` directory is
+    missing from PATH, which is how the proxy gets started under launchers and
+    init systems. Any other executable name is left untouched.
+    """
+    if not argv or argv[0] != PRISMA_CONSOLE_SCRIPT:
+        return tuple(argv)
+    if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None:
+        return tuple(argv)
+    return (sys.executable, "-m", PRISMA_CONSOLE_SCRIPT, *argv[1:])
+
+
 def run_prisma(
     argv: Sequence[str],
     *,
@@ -200,7 +225,7 @@ def run_prisma(
     text unless ``stdout``/``stderr`` say otherwise.
     """
     with subprocess.Popen(
-        argv,
+        resolve_prisma_argv(argv),
         env=env,
         stdout=stdout,
         stderr=stderr,

litellm/proxy/prisma_migration.py

@@ -14,6 +14,8 @@ sys.path.insert(0, os.path.abspath("./"))
 
 from typing import Final
 
+from litellm_proxy_extras.prisma_toolchain import resolve_prisma_argv
+
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy.proxy_cli import run_server
 from litellm.secret_managers.main import str_to_bool
@@ -29,7 +31,7 @@ def main() -> int:
     run_server(run_server_args, standalone_mode=False)
 
     verbose_proxy_logger.info("Running 'prisma generate'...")
-    result: Final = subprocess.run(("prisma", "generate"), capture_output=True, text=True)
+    result: Final = subprocess.run(resolve_prisma_argv(("prisma", "generate")), capture_output=True, text=True)
     verbose_proxy_logger.info("'prisma generate' stdout: %s", result.stdout)
 
     if result.returncode != 0:

litellm/proxy/proxy_cli.py

@@ -1267,73 +1267,69 @@ def run_server(
                         flush=True,
                     )
                     sys.exit(1)
-            try:
-                from litellm.secret_managers.main import get_secret
+            from litellm.secret_managers.main import get_secret
 
-                connection_url_params: Final = _build_db_connection_url_params(
-                    connection_limit=db_connection_pool_limit,
-                    pool_timeout=db_connection_timeout,
-                    connect_timeout=db_connect_timeout,
-                    socket_timeout=db_socket_timeout,
-                    disable_prepared_statements=db_disable_prepared_statements,
-                    extra_params=db_extra_connection_params,
+            connection_url_params: Final = _build_db_connection_url_params(
+                connection_limit=db_connection_pool_limit,
+                pool_timeout=db_connection_timeout,
+                connect_timeout=db_connect_timeout,
+                socket_timeout=db_socket_timeout,
+                disable_prepared_statements=db_disable_prepared_statements,
+                extra_params=db_extra_connection_params,
+            )
+            lifetime_params: Final = idle_lifetime_params(general_settings.get("database_max_idle_connection_lifetime"))
+            if os.getenv("DATABASE_URL", None) is not None:
+                database_url = get_secret("DATABASE_URL", default_value=None)
+                resolved_url: Final[str | None] = str(database_url) if database_url else None
+                pg_options: Final[str] = _pg_options_with_timeouts(
+                    _url_query_value(resolved_url, "options"),
+                    db_statement_timeout,
+                    db_lock_timeout,
                 )
-                lifetime_params: Final = idle_lifetime_params(
-                    general_settings.get("database_max_idle_connection_lifetime")
+                writer_url: Final = (
+                    _with_query_value(resolved_url, "options", pg_options)
+                    if resolved_url and pg_options
+                    else resolved_url
                 )
-                if os.getenv("DATABASE_URL", None) is not None:
-                    database_url = get_secret("DATABASE_URL", default_value=None)
-                    resolved_url: Final[str | None] = str(database_url) if database_url else None
-                    pg_options: Final[str] = _pg_options_with_timeouts(
-                        _url_query_value(resolved_url, "options"),
-                        db_statement_timeout,
-                        db_lock_timeout,
-                    )
-                    writer_url: Final = (
-                        _with_query_value(resolved_url, "options", pg_options)
-                        if resolved_url and pg_options
-                        else resolved_url
-                    )
-                    modified_url = append_query_params(
-                        writer_url,
-                        connection_url_params,
-                    )
-                    os.environ["DATABASE_URL"] = translate_libpq_ssl_params(
-                        add_missing_query_params(modified_url, lifetime_params)
-                    )
-                if os.getenv("DIRECT_URL", None) is not None:
-                    database_url = os.getenv("DIRECT_URL")
-                    modified_url = append_query_params(database_url, connection_url_params)
-                    os.environ["DIRECT_URL"] = translate_libpq_ssl_params(
-                        add_missing_query_params(modified_url, lifetime_params)
-                    )
-                # The reader pool is a real pool against the same configured cap, so it
-                # gets the allowlisted pool params. Schema-affecting ones, including any
-                # the operator smuggled in through database_extra_connection_params, stay
-                # on the writer. Anything pinned on the replica URL wins, unlike the
-                # writer where the config is applied on top.
-                read_replica_url: Final[str | None] = os.getenv("DATABASE_URL_READ_REPLICA")
-                if read_replica_url:
-                    reader_options: Final[str] = _pg_options_with_timeouts(
-                        _url_query_value(read_replica_url, "options"),
-                        db_statement_timeout,
-                        db_lock_timeout,
-                    )
-                    os.environ["DATABASE_URL_READ_REPLICA"] = translate_libpq_ssl_params(
+                modified_url = append_query_params(
+                    writer_url,
+                    connection_url_params,
+                )
+                os.environ["DATABASE_URL"] = translate_libpq_ssl_params(
+                    add_missing_query_params(modified_url, lifetime_params)
+                )
+            if os.getenv("DIRECT_URL", None) is not None:
+                database_url = os.getenv("DIRECT_URL")
+                modified_url = append_query_params(database_url, connection_url_params)
+                os.environ["DIRECT_URL"] = translate_libpq_ssl_params(
+                    add_missing_query_params(modified_url, lifetime_params)
+                )
+            # The reader pool is a real pool against the same configured cap, so it
+            # gets the allowlisted pool params. Schema-affecting ones, including any
+            # the operator smuggled in through database_extra_connection_params, stay
+            # on the writer. Anything pinned on the replica URL wins, unlike the
+            # writer where the config is applied on top.
+            read_replica_url: Final[str | None] = os.getenv("DATABASE_URL_READ_REPLICA")
+            if read_replica_url:
+                reader_options: Final[str] = _pg_options_with_timeouts(
+                    _url_query_value(read_replica_url, "options"),
+                    db_statement_timeout,
+                    db_lock_timeout,
+                )
+                os.environ["DATABASE_URL_READ_REPLICA"] = translate_libpq_ssl_params(
+                    add_missing_query_params(
                         add_missing_query_params(
-                            add_missing_query_params(
-                                _with_query_value(read_replica_url, "options", reader_options)
-                                if reader_options
-                                else read_replica_url,
-                                reader_shareable_params(connection_url_params),
-                            ),
-                            lifetime_params,
-                        )
+                            _with_query_value(read_replica_url, "options", reader_options)
+                            if reader_options
+                            else read_replica_url,
+                            reader_shareable_params(connection_url_params),
+                        ),
+                        lifetime_params,
                     )
-                subprocess.run(["prisma"], capture_output=True)
-                is_prisma_runnable = True
-            except FileNotFoundError:
-                is_prisma_runnable = False
+                )
+            from litellm_proxy_extras.prisma_toolchain import prisma_cli_available
+
+            is_prisma_runnable: Final = prisma_cli_available()
 
             if is_prisma_runnable:
                 from litellm.proxy.db.check_migration import check_prisma_schema_diff
@@ -1382,7 +1378,8 @@ def run_server(
                             )
             else:
                 print(
-                    f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found."  # noqa: F541
+                    "Unable to connect to DB. DATABASE_URL found in environment, but the prisma CLI is neither on "
+                    "PATH nor importable as a package."
                 )
         pgbouncer_settings: Final = PgBouncerSettings()
         upstream_database_url: Final = os.getenv("DATABASE_URL")

tests/proxy_migration_tests/test_prisma_toolchain.py

@@ -37,7 +37,10 @@ from litellm_proxy_extras.prisma_toolchain import (
     node_binary_path,
     prisma_bootstrap_timeout,
     prisma_command_timeout,
+    prisma_cli_available,
     prisma_migrate_deploy_timeout,
+    resolve_prisma_argv,
+    run_prisma,
 )
 from litellm_proxy_extras.utils import ProxyExtrasDBManager
 
@@ -401,3 +404,95 @@ def test_every_prisma_command_timeout_is_overridable(module: str) -> None:
         f"{module} still hardcodes a Prisma timeout at lines {literals}; "
         "route it through prisma_command_timeout() so it can be raised without a release"
     )
+
+
+FAKE_PRISMA_MODULE_MAIN = """import json
+import sys
+
+print(json.dumps({"module_argv": sys.argv[1:]}))
+"""
+
+
+def _write_fake_prisma_module(tmp_path: Path) -> Path:
+    package_dir = tmp_path / "fakemodule" / "prisma"
+    package_dir.mkdir(parents=True)
+    (package_dir / "__init__.py").write_text("")
+    (package_dir / "__main__.py").write_text(FAKE_PRISMA_MODULE_MAIN)
+    return package_dir.parent
+
+
+def _empty_bin(tmp_path: Path) -> Path:
+    bin_dir = tmp_path / "emptybin"
+    bin_dir.mkdir()
+    return bin_dir
+
+
+def test_run_prisma_uses_the_module_when_the_console_script_is_not_on_path(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    empty_bin = _empty_bin(tmp_path)
+    module_root = _write_fake_prisma_module(tmp_path)
+    monkeypatch.setenv("PATH", str(empty_bin))
+
+    result = run_prisma(
+        ["prisma", "migrate", "deploy"],
+        timeout=60,
+        env={"PATH": str(empty_bin), "PYTHONPATH": str(module_root)},
+    )
+
+    assert json.loads(result.stdout) == {"module_argv": ["migrate", "deploy"]}
+
+
+def test_run_prisma_prefers_the_console_script_on_path(
+    toolchain_env: tuple[Path, Path], tmp_path: Path
+) -> None:
+    _, log_path = toolchain_env
+    module_root = _write_fake_prisma_module(tmp_path)
+
+    result = run_prisma(
+        ["prisma", "--version"],
+        timeout=60,
+        env={**os.environ, "PYTHONPATH": str(module_root)},
+    )
+
+    assert [call["args"] for call in _fake_prisma_calls(log_path)] == [["--version"]]
+    assert "module_argv" not in result.stdout
+
+
+def test_resolve_prisma_argv_leaves_an_explicit_cli_path_alone(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("PATH", str(_empty_bin(tmp_path)))
+    explicit = ("/app/.cache/prisma-python/prisma", "migrate", "deploy")
+
+    assert resolve_prisma_argv(explicit) == explicit
+
+
+def test_prisma_cli_is_unavailable_with_neither_script_nor_package(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("PATH", str(_empty_bin(tmp_path)))
+    monkeypatch.delitem(sys.modules, "prisma", raising=False)
+    monkeypatch.setattr(sys, "path", [])
+
+    assert prisma_cli_available() is False
+
+
+def test_prisma_cli_is_available_through_the_package_alone(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("PATH", str(_empty_bin(tmp_path)))
+    monkeypatch.delitem(sys.modules, "prisma", raising=False)
+    monkeypatch.setattr(sys, "path", [str(_write_fake_prisma_module(tmp_path))])
+
+    assert prisma_cli_available() is True
+
+
+def test_prisma_cli_is_available_through_the_console_script_alone(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("PATH", str(_write_fake_prisma(tmp_path)))
+    monkeypatch.delitem(sys.modules, "prisma", raising=False)
+    monkeypatch.setattr(sys, "path", [])
+
+    assert prisma_cli_available() is True

tests/test_litellm/proxy/test_prisma_migration.py

@@ -1,4 +1,6 @@
 import os
+import sys
+from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -64,3 +66,34 @@ class TestPrismaMigration:
                 prisma_migration.main()
 
         mock_subprocess_run.assert_not_called()
+
+    @patch("litellm.proxy.prisma_migration.subprocess.run")  # test-quality-ok: the spawned argv is the behavior under test
+    @patch("litellm.proxy.prisma_migration.run_server")  # test-quality-ok: run_server boots the whole proxy
+    def test_prisma_generate_runs_through_the_module_when_the_cli_is_not_on_path(
+        self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, tmp_path: Path
+    ) -> None:
+        mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="")
+        empty_bin: Path = tmp_path / "emptybin"
+        empty_bin.mkdir()
+
+        with patch.dict(os.environ, {"PATH": str(empty_bin)}, clear=True):
+            assert prisma_migration.main() == 0
+
+        assert mock_subprocess_run.call_args.args[0] == (sys.executable, "-m", "prisma", "generate")
+
+    @patch("litellm.proxy.prisma_migration.subprocess.run")  # test-quality-ok: the spawned argv is the behavior under test
+    @patch("litellm.proxy.prisma_migration.run_server")  # test-quality-ok: run_server boots the whole proxy
+    def test_prisma_generate_runs_the_console_script_when_it_is_on_path(
+        self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, tmp_path: Path
+    ) -> None:
+        mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="")
+        bin_dir: Path = tmp_path / "bin"
+        bin_dir.mkdir()
+        script: Path = bin_dir / "prisma"
+        script.write_text("#!/bin/sh\nexit 0\n")
+        script.chmod(0o755)
+
+        with patch.dict(os.environ, {"PATH": str(bin_dir)}, clear=True):
+            assert prisma_migration.main() == 0
+
+        assert mock_subprocess_run.call_args.args[0] == ("prisma", "generate")

tests/test_litellm/proxy/test_proxy_cli.py

@@ -1940,6 +1940,66 @@ class TestRunServerDbSetup:
                 use_migrate=False, use_v2_resolver=False
             )
 
+    @patch("atexit.register")
+    @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")  # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
+    @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")  # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
+    @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")  # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
+    def test_migrations_run_when_the_prisma_cli_is_not_on_path(
+        self,
+        mock_should_update_schema,
+        mock_check_schema_diff,
+        mock_setup_database,
+        mock_atexit_register,
+        tmp_path,
+        capsys,
+    ):
+        from litellm.proxy.proxy_cli import run_server
+
+        mock_should_update_schema.return_value = True
+        empty_bin = tmp_path / "emptybin"
+        empty_bin.mkdir()
+
+        mock_proxy_module = MagicMock(
+            app=MagicMock(),
+            ProxyConfig=MagicMock(),
+            KeyManagementSettings=MagicMock(),
+            save_worker_config=MagicMock(),
+        )
+
+        clean_env = {
+            k: v
+            for k, v in os.environ.items()
+            if k not in ("DATABASE_URL", "DIRECT_URL")
+        }
+        clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
+        clean_env["PATH"] = str(empty_bin)
+
+        with (
+            patch.dict(os.environ, clean_env, clear=True),
+            patch.dict(
+                "sys.modules",
+                {
+                    "proxy_server": mock_proxy_module,
+                    "litellm.proxy.proxy_server": mock_proxy_module,
+                },
+            ),
+            patch(  # test-quality-ok: same isolation as the sibling CLI tests above
+                "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
+            ) as mock_get_args,
+        ):
+            mock_get_args.return_value = {
+                "app": "litellm.proxy.proxy_server:app",
+                "host": "localhost",
+                "port": 8000,
+            }
+
+            run_server.main(["--local", "--skip_server_startup"], standalone_mode=False)
+
+        assert "prisma CLI is neither on PATH" not in capsys.readouterr().out
+        mock_setup_database.assert_called_once_with(
+            use_migrate=True, use_v2_resolver=False
+        )
+
     @patch("subprocess.run")
     @patch("atexit.register")
     @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")