fix(proxy-extras): run prisma migrate deploy under the pinned coordinator lock instead of Prisma's session advisory lock (#46368)

Through a transaction-mode pooler Prisma takes pg_advisory_lock(72707369) on one
backend and releases it on another, so the first backend keeps the lock while idle
and the next boot waits on it. Migrations now run with
PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK=1 while LiteLLM holds the coordinator lock and
Prisma's key as transaction-level locks on one pinned connection, so nothing can
outlive the migration

Resolves LIT-9571

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm-proxy-extras/litellm_proxy_extras/migration_lock.py

@@ -13,6 +13,7 @@ from litellm_proxy_extras.prisma_toolchain import (
 )
 
 MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
+PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
 
 if TYPE_CHECKING:
     import psycopg
@@ -32,7 +33,7 @@ def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
     if not database_url:
         return environment
     migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL"))
-    return {**environment, "DATABASE_URL": migration_url}
+    return {**environment, "DATABASE_URL": migration_url, "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK": "1"}
 
 
 @dataclass(frozen=True, slots=True)
@@ -55,10 +56,15 @@ class MigrationCoordinator:
     def check_connection(self) -> None:
         self.connection.execute("SELECT 1")
 
+    def release_statement_snapshot(self) -> None:
+        # a parameterized statement leaves its portal snapshot registered on the backend until the
+        # next statement, and CREATE INDEX CONCURRENTLY waits for every such snapshot to go away
+        self.connection.execute("SELECT 1")
+
     def acquire_prisma_lock(self) -> None:
         deadline: Final = time.monotonic() + migration_lock_timeout()
         while time.monotonic() < deadline:
-            if _try_lock(self.connection, 72707369):
+            if _try_lock(self.connection, PRISMA_ADVISORY_LOCK_KEY):
                 return
             time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
         raise RuntimeError(
@@ -95,6 +101,17 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
     )
 
 
+@contextmanager
+def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
+    """Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one
+    pinned connection while a Prisma command runs with its session advisory lock disabled, so a
+    transaction pooler can never strand that lock on an idle backend after the command exits."""
+    with migration_lock(database_url) as coordinator:
+        coordinator.acquire_prisma_lock()
+        coordinator.release_statement_snapshot()
+        yield
+
+
 @contextmanager
 def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]:
     """A session-level, non-blocking hold of the migration coordinator lock on an autocommit

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -9,6 +9,7 @@ import tempfile
 import threading
 import time
 from collections.abc import Callable
+from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
@@ -1021,7 +1022,11 @@ class ProxyExtrasDBManager:
     def _setup_database_v2(use_migrate: bool) -> bool:
         if not use_migrate:
             return ProxyExtrasDBManager._run_database_v2(False)
-        from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
+        from litellm_proxy_extras.migration_lock import (
+            migration_environment,
+            migration_lock,
+            prisma_migration_lock,
+        )
         from litellm_proxy_extras.migration_recovery import (
             baseline_current_schema,
             recover_completed_migration,
@@ -1055,7 +1060,9 @@ class ProxyExtrasDBManager:
                     migration_environment(_get_prisma_env()),
                 )
 
-        while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing):
+        while not ProxyExtrasDBManager._run_database_v2(
+            True, recover_completed, baseline_existing, lambda: prisma_migration_lock(lock_url)
+        ):
             continue
         return True
 
@@ -1064,6 +1071,7 @@ class ProxyExtrasDBManager:
         use_migrate: bool,
         recover_completed: Callable[[str], bool] = lambda name: False,
         baseline_existing: "Callable[[str], None] | None" = None,
+        deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
     ) -> bool:
         """
         v2 migration resolver (opt-in via --use_v2_migration_resolver).
@@ -1126,11 +1134,12 @@ class ProxyExtrasDBManager:
             while not budget.exhausted:
                 attempt_started = time.monotonic()
                 try:
-                    result = prisma_toolchain.run_prisma(
-                        [_get_prisma_command(), "migrate", "deploy"],
-                        timeout=deploy_timeout,
-                        env=migration_env,
-                    )
+                    with deploy_lock():
+                        result = prisma_toolchain.run_prisma(
+                            [_get_prisma_command(), "migrate", "deploy"],
+                            timeout=deploy_timeout,
+                            env=migration_env,
+                        )
                     logger.info(f"prisma migrate deploy stdout: {result.stdout}")
                     return True
 
@@ -1155,6 +1164,7 @@ class ProxyExtrasDBManager:
                         e,
                         budget,
                         time.monotonic() - attempt_started,
+                        deploy_lock=deploy_lock,
                     )
 
                 if next_budget.attempts_left < budget.attempts_left:
@@ -1187,15 +1197,19 @@ class ProxyExtrasDBManager:
         return match.group(1) if match else None
 
     @staticmethod
-    def _v2_roll_back_migration_best_effort(migration_name: str) -> None:
+    def _v2_roll_back_migration_best_effort(
+        migration_name: str,
+        deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
+    ) -> None:
         from litellm_proxy_extras.migration_lock import migration_environment
 
         try:
-            prisma_toolchain.run_prisma(
-                [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
-                timeout=prisma_command_timeout(),
-                env=migration_environment(_get_prisma_env()),
-            )
+            with deploy_lock():
+                prisma_toolchain.run_prisma(
+                    [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
+                    timeout=prisma_command_timeout(),
+                    env=migration_environment(_get_prisma_env()),
+                )
         except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
             pass
 
@@ -1204,6 +1218,7 @@ class ProxyExtrasDBManager:
         error: subprocess.CalledProcessError,
         budget: "_MigrateAttemptBudget",
         attempt_seconds: float = 0.0,
+        deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
     ) -> "_MigrateAttemptBudget":
         """Recover from one failed `prisma migrate deploy`, and price the pass.
 
@@ -1225,14 +1240,14 @@ class ProxyExtrasDBManager:
                         "deadlock race, rolling its ledger row back and retrying",
                         migration_name,
                     )
-                    ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
+                    ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
                     return budget.spend()
                 if ledger_logs and _MIGRATION_LOCK_TIMEOUT_MARKER in ledger_logs:
                     logger.info(
                         "Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying",
                         migration_name,
                     )
-                    ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
+                    ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
                     return budget.spend()
                 if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at):
                     logger.info(
@@ -1276,7 +1291,7 @@ class ProxyExtrasDBManager:
                     "rolling its ledger row back and retrying",
                     migration_name,
                 )
-                ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
+                ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
                 return budget.spend()
 
             if migration_name and _MIGRATION_LOCK_TIMEOUT_MARKER in stderr:
@@ -1286,7 +1301,7 @@ class ProxyExtrasDBManager:
                     "its ledger row back and retrying. Check for long-running transactions on the migrated tables.",
                     migration_name,
                 )
-                ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
+                ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
                 return budget.spend()
 
             raise RuntimeError(

tests/integration/database/test_prisma_migration_lock.py

@@ -0,0 +1,58 @@
+import os
+import secrets
+from typing import Final
+
+import psycopg
+import pytest
+from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, PRISMA_ADVISORY_LOCK_KEY, prisma_migration_lock
+
+pytestmark = pytest.mark.skipif("DATABASE_URL" not in os.environ, reason="needs a live Postgres")
+
+
+def _can_take_prismas_session_lock(connection: psycopg.Connection[tuple[object, ...]]) -> bool:
+    row: Final = connection.execute("SELECT pg_try_advisory_lock(%s)", (PRISMA_ADVISORY_LOCK_KEY,)).fetchone()
+    taken: Final = row is not None and row[0] is True
+    if taken:
+        connection.execute("SELECT pg_advisory_unlock(%s)", (PRISMA_ADVISORY_LOCK_KEY,))
+    return taken
+
+
+def _granted_advisory_keys(connection: psycopg.Connection[tuple[object, ...]]) -> frozenset[int]:
+    rows: Final = connection.execute(
+        "SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted"
+    ).fetchall()
+    return frozenset(int(row[0]) for row in rows)
+
+
+def _index_builds_concurrently_now(connection: psycopg.Connection[tuple[object, ...]]) -> bool:
+    table: Final = f'"lit_cic_{secrets.token_hex(4)}"'
+    connection.execute(f"CREATE TABLE {table} (a int)")
+    try:
+        connection.execute("SET lock_timeout = '5s'")
+        connection.execute(f"CREATE INDEX CONCURRENTLY ON {table} (a)")
+    except psycopg.errors.LockNotAvailable:
+        return False
+    finally:
+        connection.execute("RESET lock_timeout")
+        connection.execute(f"DROP TABLE {table}")
+    return True
+
+
+def test_prisma_migration_lock_does_not_stall_a_concurrent_index_build():
+    with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
+        with prisma_migration_lock(os.environ["DATABASE_URL"]):
+            assert _index_builds_concurrently_now(peer)
+
+
+def test_prisma_migration_lock_holds_prismas_key_on_a_pinned_connection_and_leaves_nothing_behind():
+    with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
+        assert _can_take_prismas_session_lock(peer)
+
+        with prisma_migration_lock(os.environ["DATABASE_URL"]):
+            held: Final = _granted_advisory_keys(peer)
+            assert PRISMA_ADVISORY_LOCK_KEY in held
+            assert MIGRATION_LOCK_KEY in held
+            assert not _can_take_prismas_session_lock(peer)
+
+        assert _granted_advisory_keys(peer).isdisjoint({PRISMA_ADVISORY_LOCK_KEY, MIGRATION_LOCK_KEY})
+        assert _can_take_prismas_session_lock(peer)

tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py

@@ -5,6 +5,7 @@ import re
 import subprocess
 import sys
 import threading
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Final, NoReturn, Optional
@@ -2083,3 +2084,101 @@ def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch
 
     assert ProxyExtrasDBManager.repair_invalid_indexes() is False
 
+
+@pytest.mark.parametrize(
+    "environment",
+    (
+        {"DATABASE_URL": "postgresql://pool/db"},
+        {"DATABASE_URL": "postgresql://pool/db?pgbouncer=true", "DIRECT_URL": "postgresql://writer/db"},
+    ),
+    ids=("pooled-only", "pooled-with-direct"),
+)
+def test_prisma_runs_its_migrations_without_a_session_advisory_lock(environment):
+    from litellm_proxy_extras.migration_lock import migration_environment
+
+    migrated = migration_environment(environment)
+
+    assert migrated["PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK"] == "1"
+    assert "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK" not in migration_environment({"PRISMA_OFFLINE_MODE": "true"})
+
+
+class TestMigrateDeployRunsUnderTheCoordinatorLock:
+    """`prisma migrate deploy` only runs while the injected lock is held, and the lock is
+    released again before the resolver handles the outcome, success or failure."""
+
+    @staticmethod
+    def _recording_lock(events: list[str]):
+        from contextlib import contextmanager
+
+        @contextmanager
+        def lock():
+            events.append("lock")
+            try:
+                yield
+            finally:
+                events.append("unlock")
+
+        return lock
+
+    def test_deploy_runs_between_lock_and_unlock(self, monkeypatch, tmp_path):
+        import litellm_proxy_extras.utils as utils_module
+
+        events: list[str] = []
+        harness = _MigrateDeployHarness(monkeypatch, tmp_path, ["ok"])
+        original_run = harness._fake_run
+
+        def run_and_record(cmd, **kwargs):
+            events.append("deploy")
+            return original_run(cmd, **kwargs)
+
+        monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", run_and_record)
+
+        assert ProxyExtrasDBManager._run_database_v2(True, deploy_lock=self._recording_lock(events)) is True
+        assert events == ["lock", "deploy", "unlock"]
+
+    def test_a_failed_deploy_releases_the_lock_before_the_retry(self, monkeypatch, tmp_path):
+        import litellm_proxy_extras.utils as utils_module
+
+        events: list[str] = []
+        harness = _MigrateDeployHarness(
+            monkeypatch, tmp_path, ["Error: P1002\nTimed out waiting for the advisory lock", "ok"]
+        )
+        original_run = harness._fake_run
+
+        def run_and_record(cmd, **kwargs):
+            events.append("deploy")
+            return original_run(cmd, **kwargs)
+
+        monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", run_and_record)
+
+        assert ProxyExtrasDBManager._run_database_v2(True, deploy_lock=self._recording_lock(events)) is True
+        assert events == ["lock", "deploy", "unlock", "lock", "deploy", "unlock"]
+
+    def test_a_deadlock_rollback_runs_under_the_coordinator_lock(
+        self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+    ) -> None:
+        import litellm_proxy_extras.utils as utils_module
+
+        migration_name: Final = "20260921190000_agent_identity"
+        stderr: Final = (
+            "Error: P3018\n\n"
+            "A migration failed to apply. New migrations cannot be applied before the error is recovered from.\n\n"
+            f"Migration name: {migration_name}\n\n"
+            "Database error code: 40P01\n\n"
+            f"Database error:\n{_P3009_DEADLOCK_LOGS}"
+        )
+        events: Final[list[str]] = []
+        harness: Final = _MigrateDeployHarness(monkeypatch, tmp_path, [stderr, "ok"])
+        original_run: Final = harness._fake_run
+
+        def run_and_record(cmd: Sequence[str], *, timeout: float, env: Mapping[str, str]) -> _FakeCompleted:
+            if cmd[1:] == ["migrate", "deploy"]:
+                events.append("deploy")
+            elif cmd[1:4] == ["migrate", "resolve", "--rolled-back"]:
+                events.append("rollback")
+            return original_run(cmd, timeout=timeout, env=env)
+
+        monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", run_and_record)
+
+        assert ProxyExtrasDBManager._run_database_v2(True, deploy_lock=self._recording_lock(events)) is True
+        assert events == ["lock", "deploy", "unlock", "lock", "rollback", "unlock", "lock", "deploy", "unlock"]