diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2062ca93fb3..245244250ee 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -353,9 +353,8 @@ class ProxyExtrasDBManager: pass @staticmethod - def _failed_migration_logs(migration_name: str) -> Optional[str]: - """Return failed migration logs, or None if the ledger is unavailable.""" - database_url = os.getenv("DATABASE_URL") + def _read_migration_ledger(query: str, params: tuple[str, ...]) -> "tuple[object, ...] | None": + database_url: Final = os.getenv("DATABASE_URL") if not database_url: return None @@ -364,28 +363,37 @@ class ProxyExtrasDBManager: except ImportError: return None - cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) - ledger_table = psycopg.sql.SQL("{}.{}").format( - psycopg.sql.Identifier( - ProxyExtrasDBManager._prisma_schema_param(database_url) or "public" - ), + cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + ledger_table: Final = psycopg.sql.SQL("{}.{}").format( + psycopg.sql.Identifier(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"), psycopg.sql.Identifier("_prisma_migrations"), ) try: - with psycopg.connect( - cleaned_url, connect_timeout=10, autocommit=True - ) as conn: - row = conn.execute( - psycopg.sql.SQL( - "SELECT logs FROM {} " - "WHERE migration_name = %s AND finished_at IS NULL " - "AND rolled_back_at IS NULL" - ).format(ledger_table), - (migration_name,), - ).fetchone() + with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: + row: Final = conn.execute(psycopg.sql.SQL(query).format(ledger_table), params).fetchone() except (psycopg.OperationalError, psycopg.DatabaseError): return None - return (row[0] or "") if row else "" + return tuple(row) if row is not None else () + + @staticmethod + def _failed_migration_logs(migration_name: str, started_at: str) -> Optional[str]: + row: Final = ProxyExtrasDBManager._read_migration_ledger( + "SELECT logs FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz " + "AND finished_at IS NULL AND rolled_back_at IS NULL", + (migration_name, started_at), + ) + if row is None: + return None + return row[0] if row and isinstance(row[0], str) else "" + + @staticmethod + def _failed_migration_recovered(migration_name: str, started_at: str) -> bool: + row: Final = ProxyExtrasDBManager._read_migration_ledger( + "SELECT 1 FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz " + "AND (finished_at IS NOT NULL OR rolled_back_at IS NOT NULL)", + (migration_name, started_at), + ) + return bool(row) @staticmethod def _resolve_specific_migration(migration_name: str): @@ -1102,6 +1110,11 @@ class ProxyExtrasDBManager: return match.group(1) if match else None return None + @staticmethod + def _v2_failed_migration_started_at(stderr: str, migration_name: str) -> "str | None": + match: Final = re.search(rf"`{re.escape(migration_name)}` migration started at ([^\r\n]+?) failed", stderr) + return match.group(1) if match else None + @staticmethod def _v2_roll_back_migration_best_effort(migration_name: str) -> None: from litellm_proxy_extras.migration_lock import migration_environment @@ -1130,8 +1143,11 @@ class ProxyExtrasDBManager: if "P3009" in stderr: migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr) - if migration_name: - ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name) + started_at: Final = ( + ProxyExtrasDBManager._v2_failed_migration_started_at(stderr, migration_name) if migration_name else None + ) + if migration_name and started_at: + ledger_logs: Final = ProxyExtrasDBManager._failed_migration_logs(migration_name, started_at) if ledger_logs and _MIGRATION_DEADLOCK_MARKER in ledger_logs: logger.info( "Migration %s failed in a concurrent migrate deploy " @@ -1140,6 +1156,14 @@ class ProxyExtrasDBManager: ) ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) return budget.spend() + if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at): + logger.info( + "Migration %s started at %s was already rolled back or completed by a concurrent " + "migrate deploy, retrying", + migration_name, + started_at, + ) + return budget.spend() raise RuntimeError( "Migration completion could not be verified. LiteLLM startup has stopped.\n\n" f"Prisma migration history (migration name and start time):\n{stderr}\n\n" diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index 29c9ec56d91..ea4a25283a1 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -3,6 +3,7 @@ import os import re import sys import threading +from dataclasses import dataclass from pathlib import Path from typing import Final @@ -709,6 +710,9 @@ class TestSpendLogsPartitionDetectionMissingPsycopg: _ATTEMPT_BUDGET = 4 +_P3009_MIGRATION_NAME = "20260415120000_health_check_latest_per_model_index" +_P3009_STARTED_AT = "2026-10-02 23:20:56.439594 UTC" +_P3009_DEADLOCK_LOGS = "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock on transaction 991" _P3005_STDERR = """Error: P3005 @@ -730,6 +734,80 @@ ERROR: relation "SomeTable" already exists """ +def _p3009_stderr(migration_name: str, started_at: str) -> str: + return ( + "Error: P3009\n\n" + "migrate found failed migrations in the target database, new migrations will not be applied. " + "Read more about how to resolve migration issues in a production database: " + "https://pris.ly/d/migrate-resolve\n" + f"The `{migration_name}` migration started at {started_at} failed\n" + ) + + +@dataclass(frozen=True, slots=True) +class _LedgerRow: + migration_name: str + started_at: str + finished: bool = False + rolled_back: bool = False + logs: str | None = None + + +@dataclass(frozen=True, slots=True) +class _LedgerCursor: + row: tuple[object, ...] | None = None + + def fetchone(self) -> tuple[object, ...] | None: + return self.row + + def fetchall(self) -> tuple[tuple[object, ...], ...]: + return () + + +class _LedgerConnection: + def __init__(self, ledger: "_FakeLedger") -> None: + self.ledger = ledger + + def __enter__(self) -> "_LedgerConnection": + return self + + def __exit__(self, *args: object) -> None: + return None + + def execute(self, query: object, params: tuple[object, ...] = ()) -> _LedgerCursor: + return self.ledger.execute(query, params) + + +class _FakeLedger: + def __init__(self, at_error: tuple[_LedgerRow, ...], after_peer: tuple[_LedgerRow, ...]) -> None: + self.rows = at_error + self.after_peer = after_peer + self._peer_observed = False + + def connect(self, *args: object, **kwargs: object) -> _LedgerConnection: + return _LedgerConnection(self) + + def execute(self, query: object, params: tuple[object, ...]) -> _LedgerCursor: + text: Final = str(query) + if "WHERE migration_name = %s" not in text or not params: + return _LedgerCursor() + if not self._peer_observed: + self.rows = self.after_peer + self._peer_observed = True + matching: Final = tuple( + row + for row in self.rows + if row.migration_name == params[0] and (len(params) == 1 or row.started_at == params[1]) + ) + if "rolled_back_at IS NULL" in text: + unresolved: Final = next((row for row in matching if not row.finished and not row.rolled_back), None) + return _LedgerCursor((unresolved.logs,) if unresolved else None) + if "IS NOT NULL" in text: + resolved: Final = next((row for row in matching if row.finished or row.rolled_back), None) + return _LedgerCursor((1,) if resolved else None) + return _LedgerCursor() + + @pytest.mark.parametrize( "pooled,direct,expected", ( @@ -759,7 +837,15 @@ class _MigrateDeployHarness: `prisma migrate deploy` outcomes, with every recovery command faked out so nothing touches a database or the packaged migrations directory.""" - def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False, confirmed_migrations=()): + def __init__( + self, + monkeypatch, + tmp_path, + outcomes, + repeat_last=False, + confirmed_migrations=(), + ledger: "_FakeLedger | None" = None, + ): import subprocess as subprocess_module import litellm_proxy_extras.utils as utils_module @@ -772,7 +858,11 @@ class _MigrateDeployHarness: self._subprocess_module = subprocess_module self.confirmed_migrations = set(confirmed_migrations) - monkeypatch.delenv("DATABASE_URL", raising=False) + if ledger is None: + monkeypatch.delenv("DATABASE_URL", raising=False) + else: + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setattr("psycopg.connect", ledger.connect) monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path)) monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run) monkeypatch.setattr(utils_module, "_get_prisma_env", lambda: {}) @@ -819,6 +909,84 @@ class _MigrateDeployHarness: return True +class TestConcurrentP3009Recovery: + @pytest.mark.parametrize( + "after_peer", + ( + (_LedgerRow(_P3009_MIGRATION_NAME, _P3009_STARTED_AT, rolled_back=True, logs=_P3009_DEADLOCK_LOGS),), + ( + _LedgerRow(_P3009_MIGRATION_NAME, _P3009_STARTED_AT, rolled_back=True, logs=_P3009_DEADLOCK_LOGS), + _LedgerRow(_P3009_MIGRATION_NAME, "2026-10-02 23:21:11.539224 UTC"), + ), + (_LedgerRow(_P3009_MIGRATION_NAME, _P3009_STARTED_AT, finished=True, logs=_P3009_DEADLOCK_LOGS),), + ), + ids=("rolled-back", "rolled-back-beside-a-fresh-in-flight-row", "finished"), + ) + def test_a_p3009_row_a_peer_already_recovered_is_retried( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + after_peer: tuple[_LedgerRow, ...], + ) -> None: + deadlocked_row: Final = _LedgerRow( + _P3009_MIGRATION_NAME, + _P3009_STARTED_AT, + logs=_P3009_DEADLOCK_LOGS, + ) + harness: Final = _MigrateDeployHarness( + monkeypatch, + tmp_path, + [_p3009_stderr(_P3009_MIGRATION_NAME, _P3009_STARTED_AT), "ok"], + ledger=_FakeLedger(at_error=(deadlocked_row,), after_peer=after_peer), + ) + + assert harness.run() is True + assert len(harness.deploy_calls) == 2 + + @pytest.mark.parametrize( + "ledger_rows", + ( + ( + _LedgerRow( + _P3009_MIGRATION_NAME, + _P3009_STARTED_AT, + logs='ERROR: syntax error at or near "SLECT"', + ), + ), + ( + _LedgerRow( + _P3009_MIGRATION_NAME, + _P3009_STARTED_AT, + logs='ERROR: syntax error at or near "SLECT"', + ), + _LedgerRow( + _P3009_MIGRATION_NAME, + "2026-10-02 23:19:40.120000 UTC", + rolled_back=True, + logs=_P3009_DEADLOCK_LOGS, + ), + ), + ), + ids=("only-row", "beside-a-recovered-earlier-attempt"), + ) + def test_an_unresolved_p3009_row_without_the_deadlock_marker_stops( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ledger_rows: tuple[_LedgerRow, ...], + ) -> None: + harness: Final = _MigrateDeployHarness( + monkeypatch, + tmp_path, + [_p3009_stderr(_P3009_MIGRATION_NAME, _P3009_STARTED_AT)], + ledger=_FakeLedger(at_error=ledger_rows, after_peer=ledger_rows), + ) + + with pytest.raises(RuntimeError, match="Migration completion could not be verified"): + harness.run() + assert len(harness.deploy_calls) == 1 + + class TestMigrateDeployAttemptAccounting: def test_a_push_created_database_finishes_bootstrapping(self, monkeypatch, tmp_path): harness = _MigrateDeployHarness(