diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index 2283814ab35..b51de9609d3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -18,10 +18,15 @@ recoverable one. constant: it grows with the number of pending migrations, so a fresh database that has to replay every migration this package ships overruns a per-command budget sized for the short bookkeeping commands, on a laptop as much as on a -slow CI runner. The Python ``prisma`` wrapper spawns Node and the schema engine -as separate children, so killing the wrapper on timeout leaves them running: -the retry then contends with that orphan for Prisma's advisory lock and cannot -finish any sooner. Migrate deploy therefore runs under its own budget. +slow CI runner. Migrate deploy therefore runs under its own budget. + +The Python ``prisma`` wrapper spawns Node, which spawns the Rust schema +engine, so killing only the wrapper on timeout leaves the engine running with +no parent: it keeps mutating the database after the proxy has given up, holds +Prisma's advisory lock so every retry and every later boot queues behind it, +and dies mid-migration once its pipes close, leaving a half-applied ledger row. +Every Prisma command therefore runs in a process group of its own, and a +timeout kills the whole group. All three budgets are overridable so an operator can widen them without a release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install, @@ -35,10 +40,12 @@ the deploy override says otherwise. import math import os import shutil +import signal import subprocess +from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path -from typing import Optional +from typing import IO, Optional, Union from litellm_proxy_extras._logging import logger @@ -167,6 +174,49 @@ def heal_incomplete_nodeenv_cache() -> bool: return True +def _kill_process_group(process: "subprocess.Popen[str]") -> None: + if os.name == "nt": + process.kill() + return + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + return + + +def run_prisma( + argv: Sequence[str], + *, + timeout: float, + env: Mapping[str, str], + stdout: Union[IO[str], int, None] = subprocess.PIPE, + stderr: Optional[int] = subprocess.PIPE, +) -> "subprocess.CompletedProcess[str]": + """Run one Prisma CLI command in its own process group, bounded by ``timeout``. + + Raises ``subprocess.TimeoutExpired`` once the budget is spent, after killing + the command together with every process it spawned, and + ``subprocess.CalledProcessError`` on a non-zero exit. Output is captured as + text unless ``stdout``/``stderr`` say otherwise. + """ + with subprocess.Popen( + argv, + env=env, + stdout=stdout, + stderr=stderr, + text=True, + start_new_session=True, + ) as process: + try: + out, err = process.communicate(timeout=timeout) + except BaseException: + _kill_process_group(process) + raise + if process.returncode: + raise subprocess.CalledProcessError(process.returncode, process.args, out, err) + return subprocess.CompletedProcess(process.args, process.returncode, out, err) + + def ensure_prisma_toolchain( prisma_command: str, prisma_env: dict[str, str] ) -> ToolchainBootstrap: @@ -179,14 +229,7 @@ def ensure_prisma_toolchain( timeout = prisma_bootstrap_timeout() logger.info("Preparing the Prisma CLI toolchain (timeout %ss)", timeout) try: - subprocess.run( - [prisma_command, BOOTSTRAP_ARG], - timeout=timeout, - check=True, - capture_output=True, - text=True, - env=prisma_env, - ) + run_prisma([prisma_command, BOOTSTRAP_ARG], timeout=timeout, env=prisma_env) except subprocess.TimeoutExpired: logger.warning( "Preparing the Prisma CLI toolchain timed out after %ss. Raise %s " diff --git a/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py b/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py index 157d595404e..3a5865a54cd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py +++ b/litellm-proxy-extras/litellm_proxy_extras/replica_identity.py @@ -16,7 +16,7 @@ import tempfile from pathlib import Path from litellm_proxy_extras._logging import logger -from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout +from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout, run_prisma REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL" @@ -66,7 +66,7 @@ def apply_replica_identity_full( with tempfile.TemporaryDirectory(prefix="litellm_replica_identity_") as tmp_dir: sql_path = Path(tmp_dir) / "replica_identity_full.sql" sql_path.write_text(REPLICA_IDENTITY_FULL_SQL) - subprocess.run( + run_prisma( [ prisma_command, "db", @@ -77,9 +77,6 @@ def apply_replica_identity_full( schema_path, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=prisma_env, ) except subprocess.CalledProcessError as e: diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index a4ea4789b49..71e7e9c683b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -10,6 +10,7 @@ from dataclasses import dataclass, replace from pathlib import Path from typing import Optional +from litellm_proxy_extras import prisma_toolchain from litellm_proxy_extras._logging import logger from litellm_proxy_extras.replica_identity import ( REPLICA_IDENTITY_FULL_ENV_VAR, @@ -231,7 +232,7 @@ class ProxyExtrasDBManager: # 1. Generate migration SQL file by comparing empty state to current db state logger.info("Generating baseline migration...") migration_file = init_dir / "migration.sql" - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -242,14 +243,13 @@ class ProxyExtrasDBManager: "--script", ], stdout=open(migration_file, "w"), - check=True, timeout=prisma_command_timeout(), env=prisma_env, ) # 3. Mark the migration as applied since it represents current state logger.info("Marking baseline migration as applied...") - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -257,7 +257,6 @@ class ProxyExtrasDBManager: "--applied", "0_init", ], - check=True, timeout=prisma_command_timeout(), env=prisma_env, ) @@ -286,7 +285,7 @@ class ProxyExtrasDBManager: """Mark a specific migration as rolled back""" # Set up environment for offline mode if configured prisma_env = _get_prisma_env() - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -295,8 +294,6 @@ class ProxyExtrasDBManager: migration_name, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, env=prisma_env, ) @@ -348,11 +345,9 @@ class ProxyExtrasDBManager: def _resolve_specific_migration(migration_name: str): """Mark a specific migration as applied""" prisma_env = _get_prisma_env() - subprocess.run( + prisma_toolchain.run_prisma( [_get_prisma_command(), "migrate", "resolve", "--applied", migration_name], timeout=prisma_command_timeout(), - check=True, - capture_output=True, env=prisma_env, ) @@ -436,7 +431,7 @@ class ProxyExtrasDBManager: try: logger.info("Generating migration diff between DB and schema.prisma...") with open(diff_sql_path, "w") as f: - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -447,7 +442,6 @@ class ProxyExtrasDBManager: schema_path, "--script", ], - check=True, timeout=prisma_command_timeout(), stdout=f, env=_get_prisma_env(), @@ -470,7 +464,7 @@ class ProxyExtrasDBManager: migration_files = sorted(Path(migrations_dir).glob("*/migration.sql")) for mig_file in migration_files: try: - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "db", @@ -481,9 +475,6 @@ class ProxyExtrasDBManager: schema_path, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info(f"Applied migration: {mig_file.parent.name}") @@ -516,7 +507,7 @@ class ProxyExtrasDBManager: applied_ok = False try: logger.info("Running prisma db execute to apply the migration diff...") - result = subprocess.run( + result = prisma_toolchain.run_prisma( [ _get_prisma_command(), "db", @@ -527,9 +518,6 @@ class ProxyExtrasDBManager: schema_path, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info(f"prisma db execute stdout: {result.stdout}") @@ -558,7 +546,7 @@ class ProxyExtrasDBManager: for migration_name in migration_names: try: logger.info(f"Resolving migration: {migration_name}") - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -567,9 +555,6 @@ class ProxyExtrasDBManager: migration_name, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.debug(f"Resolved migration: {migration_name}") @@ -762,11 +747,12 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) try: - subprocess.run( + prisma_toolchain.run_prisma( [_get_prisma_command(), "db", "push", "--accept-data-loss"], timeout=prisma_command_timeout(), - check=True, env=_get_prisma_env(), + stdout=None, + stderr=None, ) return True except ( @@ -789,12 +775,9 @@ class ProxyExtrasDBManager: try: while not budget.exhausted: try: - result = subprocess.run( + result = prisma_toolchain.run_prisma( [_get_prisma_command(), "migrate", "deploy"], timeout=deploy_timeout, - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -1031,12 +1014,9 @@ class ProxyExtrasDBManager: logger.info("Running prisma migrate deploy") try: # Set migrations directory for Prisma - result = subprocess.run( + result = prisma_toolchain.run_prisma( [_get_prisma_command(), "migrate", "deploy"], timeout=prisma_migrate_deploy_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -1108,7 +1088,7 @@ class ProxyExtrasDBManager: f"Found failed migration: {failed_migration}, marking as rolled back" ) # Mark the failed migration as rolled back - subprocess.run( + prisma_toolchain.run_prisma( [ _get_prisma_command(), "migrate", @@ -1117,9 +1097,6 @@ class ProxyExtrasDBManager: failed_migration, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info( @@ -1244,10 +1221,12 @@ class ProxyExtrasDBManager: if ProxyExtrasDBManager.spend_logs_is_partitioned(): raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR) # Use prisma db push with increased timeout - subprocess.run( + prisma_toolchain.run_prisma( [_get_prisma_command(), "db", "push", "--accept-data-loss"], timeout=prisma_command_timeout(), - check=True, + stdout=None, + stderr=None, + env=_get_prisma_env(), ) return True except subprocess.TimeoutExpired: diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 406f07eb792..040d67d25e4 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -42,7 +42,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path): "Error: P3018\nMigration name: 20250326162113_baseline\n" "Database error code: 42501\npermission denied for schema public" ) - with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises(RuntimeError, match="permission"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -60,7 +60,7 @@ def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path): "Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n" 'Reason: syntax error at or near "BRKN" LINE 42' ) - with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises(RuntimeError, match="cannot be auto-recovered"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -124,7 +124,7 @@ def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path): def fake_resolve(*args, **kwargs): resolve_called["n"] += 1 - monkeypatch.setattr("subprocess.run", fake_run) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", fake_run) monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve) ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set @@ -139,7 +139,7 @@ def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_pat (tmp_path / "schema.prisma").write_text("// stub") stderr = "db push error" - with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises(RuntimeError, match="prisma db push failed"): ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) @@ -209,7 +209,7 @@ def test_v2_resolve_specific_migration_failure_raises_runtime_error( "Error: P3009\nMigration `20260101000000_some_migration` failed\n" "relation already exists" ) - with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises( RuntimeError, match="Failed to mark migration .* as applied" ): @@ -228,7 +228,7 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): stdout = "Applied migration.\n" stderr = "" - monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult()) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", lambda *a, **kw: FakeResult()) resolve_called = {"n": 0} monkeypatch.setattr( @@ -296,7 +296,7 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path): "_resolve_specific_migration", lambda name: pytest.fail("a deadlocked migration must never be marked applied"), ) - monkeypatch.setattr("subprocess.run", _succeed_after(1, _DEADLOCK_P3018_STDERR)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, _DEADLOCK_P3018_STDERR)) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True @@ -309,7 +309,7 @@ def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path): monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None) with patch( - "subprocess.run", + "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, _DEADLOCK_P3018_STDERR), ): with pytest.raises(RuntimeError, match="after 4 attempts"): @@ -343,7 +343,7 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ "_resolve_specific_migration", lambda name: pytest.fail("a deadlocked migration must never be marked applied"), ) - monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True @@ -372,7 +372,7 @@ def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path "_resolve_specific_migration", lambda name: pytest.fail("a deadlocked migration must never be marked applied"), ) - monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True @@ -395,7 +395,7 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path): "_roll_back_migration", lambda name: pytest.fail("an unreadable ledger must not trigger a retry"), ) - monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) with pytest.raises(RuntimeError, match="cannot be auto-recovered"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -417,7 +417,7 @@ def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path): lambda name: 'ERROR: syntax error at or near "BRKN"', ) - with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)): with pytest.raises(RuntimeError, match="cannot be auto-recovered"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -427,7 +427,7 @@ def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path): waiter as victim) is retried, not fatal.""" _stub_v2_env(monkeypatch, tmp_path) monkeypatch.setattr( - "subprocess.run", _succeed_after(1, "Database error: deadlock detected") + "litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, "Database error: deadlock detected") ) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) @@ -446,7 +446,10 @@ def test_v2_advisory_lock_timeout_retries(monkeypatch, tmp_path): """v2: the advisory-lock waiter that times out while a peer's retry holds the lock retries instead of dying.""" _stub_v2_env(monkeypatch, tmp_path) - monkeypatch.setattr("subprocess.run", _succeed_after(2, _P1002_ADVISORY_LOCK_STDERR)) + monkeypatch.setattr( + "litellm_proxy_extras.prisma_toolchain.run_prisma", + _succeed_after(2, _P1002_ADVISORY_LOCK_STDERR), + ) ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True @@ -456,7 +459,7 @@ def test_v2_p1002_without_advisory_lock_context_still_raises(monkeypatch, tmp_pa """v2: a plain P1002 (database unreachable) stays fatal.""" _stub_v2_env(monkeypatch, tmp_path) stderr = "Error: P1002\n\nThe database server at `db`:`5432` was reached but timed out." - monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr)) + monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr)) with pytest.raises(RuntimeError, match="cannot be auto-recovered"): ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 22558faedad..9ffb57924b6 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -567,7 +567,7 @@ class TestResolveAllMigrationsLedger: return _FakeCompleted() return _FakeCompleted() - monkeypatch.setattr(utils_module.subprocess, "run", fake_run) + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", fake_run) ProxyExtrasDBManager._resolve_all_migrations(str(tmp_path), "schema.prisma") return calls @@ -604,9 +604,9 @@ class TestPartitionedSpendLogsPushGuard: import litellm_proxy_extras.utils as utils_module def fail_run(cmd, **kwargs): - raise AssertionError(f"subprocess.run should not be called, got: {cmd}") + raise AssertionError(f"run_prisma should not be called, got: {cmd}") - monkeypatch.setattr(utils_module.subprocess, "run", fail_run) + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", fail_run) def test_v1_db_push_fails_fast_with_guidance(self, monkeypatch): monkeypatch.setattr( diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index 733870f3239..0ed33193a9b 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -18,6 +18,7 @@ import ast import json import logging import os +import signal import sys import time from collections.abc import Callable @@ -47,6 +48,7 @@ FAKE_PRISMA = """#!{python} import json import os import pathlib +import subprocess import sys import time @@ -66,6 +68,9 @@ with log_path.open("a") as log: time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0"))) if args[:2] == ["migrate", "deploy"]: if earlier_same_command == 0: + if os.environ.get("FAKE_PRISMA_GRANDCHILD_PIDFILE"): + grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"]) + pathlib.Path(os.environ["FAKE_PRISMA_GRANDCHILD_PIDFILE"]).write_text(str(grandchild.pid)) time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "0"))) elif os.environ.get("FAKE_PRISMA_LATER_DEPLOY_STDERR"): print(os.environ["FAKE_PRISMA_LATER_DEPLOY_STDERR"], file=sys.stderr) @@ -272,6 +277,43 @@ def test_migrate_deploy_stops_at_its_own_timeout( assert elapsed < 30 +def _process_is_gone(pid: int, within_seconds: float) -> bool: + deadline = time.monotonic() + within_seconds + while time.monotonic() < deadline: + try: + os.kill(pid, 0) + except ProcessLookupError: + return True + time.sleep(0.05) + return False + + +def test_a_timed_out_migrate_deploy_takes_its_process_tree_with_it( + toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """The real CLI forks Node and a schema engine; a timeout must not leave them running.""" + _, log_path = toolchain_env + pidfile = tmp_path / "grandchild.pid" + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "60") + monkeypatch.setenv("FAKE_PRISMA_LATER_DEPLOY_STDERR", "Error: P3018 permission denied for schema public") + monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(pidfile)) + + with pytest.raises(RuntimeError, match="insufficient permissions"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + grandchild_pid = int(pidfile.read_text()) + try: + assert len(_deploy_calls(log_path)) == 2 + assert _process_is_gone(grandchild_pid, within_seconds=5) + finally: + try: + os.kill(grandchild_pid, signal.SIGKILL) + except ProcessLookupError: + pass + + def test_db_push_timeout_hint_names_the_per_command_budget( toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: diff --git a/tests/test_litellm/proxy/db/test_replica_identity.py b/tests/test_litellm/proxy/db/test_replica_identity.py index ecfc6433ab1..9738fc9bd98 100644 --- a/tests/test_litellm/proxy/db/test_replica_identity.py +++ b/tests/test_litellm/proxy/db/test_replica_identity.py @@ -29,7 +29,7 @@ def test_hands_the_alter_statement_to_the_prisma_cli(): return subprocess.CompletedProcess(cmd, 0) with patch( - "litellm_proxy_extras.replica_identity.subprocess.run", side_effect=capture + "litellm_proxy_extras.replica_identity.run_prisma", side_effect=capture ): applied = apply_replica_identity_full( schema_path="/somewhere/schema.prisma", @@ -60,7 +60,7 @@ def test_hands_the_alter_statement_to_the_prisma_cli(): ) def test_every_failure_is_reported_instead_of_raised(failure): with patch( - "litellm_proxy_extras.replica_identity.subprocess.run", side_effect=failure + "litellm_proxy_extras.replica_identity.run_prisma", side_effect=failure ): assert ( apply_replica_identity_full(