From 1239a4de428fc4ecab8170856848419268125b22 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 12:50:07 -0700 Subject: [PATCH] fix(proxy-extras): drop bogus MIGRATE_LOCK_TIMEOUT, tighten retry test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Greptile review correctly flagged that MIGRATE_LOCK_TIMEOUT is not a recognised Prisma env var — the advisory-lock timeout is hardcoded at 10s and not configurable (https://www.prisma.io/docs/orm/prisma-migrate/workflows/development-and-production#advisory-locking). The previous code set it on os.environ on every P1002 retry, which (a) did nothing inside Prisma and (b) leaked process state. - Drop the os.environ mutation. The retry mechanism that actually helps is the back-off sleep (gives an orphan lock time to be reaped by the pooler, or a peer migration time to finish). Comment + log line rewritten to describe the real mechanism. - Append PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK as a documented last-resort bypass in the terminal RuntimeError message, with the safety caveat that concurrent deploys without the lock corrupt the ledger. - Tighten test_v2_p1002_retries_then_succeeds: capture the env passed to subprocess.run on each attempt and assert DATABASE_URL is intact on both, instead of the previous (now-removed) os.environ assertion. This catches future refactors that move env construction outside the retry loop. - Drop the now-unused `import os` from the test file. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../litellm_proxy_extras/utils.py | 19 ++++++++++----- .../tests/test_setup_database_fail_fast.py | 23 ++++++++++++++----- 2 files changed, 30 insertions(+), 12 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index e42a44ba7a9..dcdd46a5f7c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -586,16 +586,18 @@ class ProxyExtrasDBManager: if "P1002" in stderr and "advisory lock" in stderr: last_was_advisory_lock = True + # Prisma's advisory-lock timeout is hardcoded at 10s + # (see https://www.prisma.io/docs/orm/prisma-migrate/workflows/development-and-production#advisory-locking). + # The mechanism that actually helps is the back-off + # sleep: it gives a stale orphan lock time to be + # reaped by the pooler, or a peer migration time to + # finish and release. logger.warning( "Advisory-lock contention on attempt %d " - "(Prisma key 72707369). Retrying with longer wait.", + "(Prisma key 72707369). Backing off before retry.", attempt + 1, ) time.sleep(random.randrange(5, 15)) - # Lengthen Prisma's lock wait on subsequent attempts. - # Picked up by _get_prisma_env()'s os.environ.copy() - # on the next iteration. - os.environ["MIGRATE_LOCK_TIMEOUT"] = "30" continue last_was_advisory_lock = False @@ -714,7 +716,12 @@ class ProxyExtrasDBManager: "(e.g. Neon `-pooler`, Supabase pgbouncer, RDS Proxy), " "set DIRECT_URL to a non-pooled URL — pooled connections " "in transaction mode can orphan session-scoped advisory " - "locks. v2 will use DIRECT_URL automatically when set." + "locks. v2 will use DIRECT_URL automatically when set.\n\n" + "Last resort: set PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK=true " + "to bypass advisory locking entirely. Only safe when " + "concurrent `prisma migrate deploy` runs are otherwise " + "serialized; two pods racing without the lock will " + "corrupt the migration ledger." ) raise RuntimeError( "Database migration failed after 4 attempts (retry loop " 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 e26642b1ccb..7c426d57ae6 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -5,7 +5,6 @@ The v2 resolver is opt-in via `--use_v2_migration_resolver` / the (default) behavior is unchanged from pre-fix. """ -import os import subprocess from unittest.mock import patch @@ -306,10 +305,12 @@ def test_v2_no_direct_url_passthrough(monkeypatch, tmp_path): def test_v2_p1002_retries_then_succeeds(monkeypatch, tmp_path): - """v2: P1002 advisory-lock contention is retried, MIGRATE_LOCK_TIMEOUT - is bumped to 30s on subsequent attempts, and a later success returns True.""" + """v2: P1002 advisory-lock contention is caught, the loop sleeps and + retries, and a later success returns True. The retry mechanism is the + sleep itself — Prisma's 10s advisory-lock timeout is not configurable — + so we verify (1) a retry occurred and (2) the second-attempt subprocess + env still carries DATABASE_URL (no env corruption between attempts).""" monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") - monkeypatch.delenv("MIGRATE_LOCK_TIMEOUT", raising=False) monkeypatch.setattr( ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None ) @@ -324,13 +325,15 @@ def test_v2_p1002_retries_then_succeeds(monkeypatch, tmp_path): ) calls = {"n": 0} + envs_seen = [] class FakeResult: stdout = "No pending migrations to apply\n" stderr = "" - def fake_run(cmd, *args, **kwargs): + def fake_run(cmd, *args, env=None, **kwargs): calls["n"] += 1 + envs_seen.append(env or {}) if calls["n"] == 1: raise subprocess.CalledProcessError( returncode=1, @@ -345,7 +348,14 @@ def test_v2_p1002_retries_then_succeeds(monkeypatch, tmp_path): ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True assert calls["n"] == 2, "should retry once after the lock contention" - assert os.environ.get("MIGRATE_LOCK_TIMEOUT") == "30" + # Both attempts must have received the deploy env with DATABASE_URL set — + # asserting on the env passed to subprocess.run (not on os.environ) so + # that future refactors which compute the env outside the loop are still + # caught. + assert len(envs_seen) == 2 + assert all( + e.get("DATABASE_URL") == "postgresql://u:p@localhost:9/x" for e in envs_seen + ) def test_v2_p1002_terminal_raises_with_remediation_hint(monkeypatch, tmp_path): @@ -377,3 +387,4 @@ def test_v2_p1002_terminal_raises_with_remediation_hint(monkeypatch, tmp_path): assert "pg_terminate_backend" in msg assert "objid = 72707369" in msg assert "DIRECT_URL" in msg + assert "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK" in msg