diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index b51de9609d3..9cd48fcf11a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/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, diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 1b95d24c011..7e3aff75cef 100644 --- a/litellm/proxy/prisma_migration.py +++ b/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: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e245367b1b4..d9045c57b41 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1260,73 +1260,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 @@ -1375,7 +1371,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." ) if port == 4000 and ProxyInitializationHelpers._is_port_in_use(port): port = random.randint(1024, 49152) diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index 0ed33193a9b..4e2274cc582 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/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 diff --git a/tests/test_litellm/proxy/test_prisma_migration.py b/tests/test_litellm/proxy/test_prisma_migration.py index 729adcfb9e0..3fc69b34213 100644 --- a/tests/test_litellm/proxy/test_prisma_migration.py +++ b/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") diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index c76ff189a8a..e25e6a59884 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/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")