From 105f99cea214b615888e4c9046502a42395aa077 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:29:55 -0700 Subject: [PATCH 1/2] fix(proxy-extras): kill the whole Prisma process group when a command times out Every Prisma CLI call now goes through one runner that starts the command in its own session and SIGKILLs the process group on timeout, so the Node process and the Rust schema engine die together with the Python wrapper instead of being reparented to pid 1, where they kept applying migrations after the proxy had given up and held the Prisma advisory lock against every retry and every later boot. Tests that faked subprocess.run now fake the runner, and the fake Prisma CLI in the migration tests forks a grandchild that must not outlive a timed-out migrate deploy. --- .../litellm_proxy_extras/prisma_toolchain.py | 69 +++++++++++++++---- .../litellm_proxy_extras/replica_identity.py | 7 +- .../litellm_proxy_extras/utils.py | 59 +++++----------- .../tests/test_setup_database_fail_fast.py | 30 ++++---- .../test_litellm_proxy_extras_utils.py | 6 +- .../test_prisma_toolchain.py | 42 +++++++++++ .../proxy/db/test_replica_identity.py | 4 +- 7 files changed, 139 insertions(+), 78 deletions(-) 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 d22484bc0e8..168c3febae2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -9,6 +9,7 @@ import time 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, @@ -198,7 +199,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", @@ -209,14 +210,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", @@ -224,7 +224,6 @@ class ProxyExtrasDBManager: "--applied", "0_init", ], - check=True, timeout=prisma_command_timeout(), env=prisma_env, ) @@ -253,7 +252,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", @@ -262,8 +261,6 @@ class ProxyExtrasDBManager: migration_name, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, env=prisma_env, ) @@ -315,11 +312,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, ) @@ -403,7 +398,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", @@ -414,7 +409,6 @@ class ProxyExtrasDBManager: schema_path, "--script", ], - check=True, timeout=prisma_command_timeout(), stdout=f, env=_get_prisma_env(), @@ -437,7 +431,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", @@ -448,9 +442,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}") @@ -483,7 +474,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", @@ -494,9 +485,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}") @@ -525,7 +513,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", @@ -534,9 +522,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}") @@ -726,11 +711,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 ( @@ -752,12 +738,9 @@ class ProxyExtrasDBManager: try: for attempt in range(4): 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}") @@ -1007,12 +990,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}") @@ -1084,7 +1064,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", @@ -1093,9 +1073,6 @@ class ProxyExtrasDBManager: failed_migration, ], timeout=prisma_command_timeout(), - check=True, - capture_output=True, - text=True, env=_get_prisma_env(), ) logger.info( @@ -1220,10 +1197,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..2fea48a57da 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,7 @@ 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 +456,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 b3d457707b8..7917ec8c00f 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( From f0f5f1ec78c5fca4272dca3ec66bf50495fd73ab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:03:17 -0700 Subject: [PATCH 2/2] test(proxy-extras): wrap an over-long monkeypatch line --- litellm-proxy-extras/tests/test_setup_database_fail_fast.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 2fea48a57da..040d67d25e4 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -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("litellm_proxy_extras.prisma_toolchain.run_prisma", _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