diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 369b6561931..e42a44ba7a9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -491,12 +491,15 @@ class ProxyExtrasDBManager: sorted_hostile = sorted(hostile) logger.warning( - "Database has %d migration(s) applied that are NEWER than any " - "migration this LiteLLM version ships. This usually means the " - "database was migrated by a newer LiteLLM deployment. Some API " - "endpoints may fail because this proxy's Prisma client does not " - "know about those schema changes. Consider upgrading this " - "deployment. Unknown: %s", + "Database has %d migration(s) applied that this build doesn't " + "ship: %s. If the names end with `_baseline_diff` these are " + "usually developer-named diff captures from a different feature " + "branch and are safe to ignore (the schema state they describe " + "is either already merged into this build under a different " + "filename, or extra harmless tables). If a non-`_baseline_diff` " + "migration appears here, the DB was likely migrated by a newer " + "LiteLLM deployment — upgrade this proxy or expect endpoints " + "touching the new schema to fail.", len(hostile), ", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""), ) @@ -546,8 +549,19 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) + last_was_advisory_lock = False try: for attempt in range(4): + # Prefer DIRECT_URL when set, mirroring _resolve_all_migrations. + # Pooler URLs (e.g. Neon `-pooler`, Supabase pgbouncer) run in + # transaction mode, which orphans the session-scoped advisory + # lock that `prisma migrate deploy` acquires — using a direct + # connection sidesteps the lock-orphaning class of failures. + deploy_env = _get_prisma_env() + direct_url = os.getenv("DIRECT_URL") + if direct_url: + deploy_env["DATABASE_URL"] = direct_url + try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], @@ -555,7 +569,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=deploy_env, ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") return True @@ -570,6 +584,22 @@ class ProxyExtrasDBManager: except subprocess.CalledProcessError as e: stderr = e.stderr or "" + if "P1002" in stderr and "advisory lock" in stderr: + last_was_advisory_lock = True + logger.warning( + "Advisory-lock contention on attempt %d " + "(Prisma key 72707369). Retrying with longer wait.", + 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 + if "P3005" in stderr and "database schema is not empty" in stderr: logger.info( "Schema exists but no migrations ledger — creating baseline" @@ -667,6 +697,25 @@ class ProxyExtrasDBManager: f"Manual intervention required.\n\nPrisma error:\n{stderr}" ) from e + if last_was_advisory_lock: + raise RuntimeError( + "Advisory lock 72707369 could not be acquired after 4 " + "attempts. If only one proxy is starting, this is almost " + "certainly a stale lock from a prior `prisma migrate " + "deploy` that was killed before releasing it.\n\n" + "Inspect lock holders:\n" + " SELECT pid, mode, granted FROM pg_locks " + "WHERE locktype = 'advisory' AND objid = 72707369;\n\n" + "Clear the orphan lock (terminates the connection " + "holding it):\n" + " SELECT pg_terminate_backend(pid) FROM pg_locks " + "WHERE locktype = 'advisory' AND objid = 72707369;\n\n" + "If your DATABASE_URL points at a connection pooler " + "(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." + ) raise RuntimeError( "Database migration failed after 4 attempts (retry loop " "exhausted by timeouts or repeated idempotent-recovery " 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 8d66bf872de..e26642b1ccb 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -5,6 +5,7 @@ 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 @@ -240,3 +241,139 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) assert ok is True assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery" + + +def test_v2_uses_direct_url_when_set(monkeypatch, tmp_path): + """v2: `prisma migrate deploy` must use DIRECT_URL when set, to sidestep + pooler-mode advisory-lock orphaning (Neon pooler, Supabase pgbouncer, + RDS Proxy, etc.).""" + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@pooler.x.neon.tech:5432/db") + monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@direct.x.neon.tech:5432/db") + monkeypatch.setattr( + ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + captured_env = {} + + class FakeResult: + stdout = "No pending migrations to apply\n" + stderr = "" + + def fake_run(cmd, *args, env=None, **kwargs): + captured_env.update(env or {}) + return FakeResult() + + monkeypatch.setattr("subprocess.run", fake_run) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + assert ( + captured_env.get("DATABASE_URL") + == "postgresql://u:p@direct.x.neon.tech:5432/db" + ) + + +def test_v2_no_direct_url_passthrough(monkeypatch, tmp_path): + """v2: when DIRECT_URL is unset, DATABASE_URL is passed through unchanged.""" + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@pooler.x.neon.tech:5432/db") + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setattr( + ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + captured_env = {} + + class FakeResult: + stdout = "No pending migrations to apply\n" + stderr = "" + + def fake_run(cmd, *args, env=None, **kwargs): + captured_env.update(env or {}) + return FakeResult() + + monkeypatch.setattr("subprocess.run", fake_run) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + assert ok is True + assert ( + captured_env.get("DATABASE_URL") + == "postgresql://u:p@pooler.x.neon.tech:5432/db" + ) + + +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.""" + 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 + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setattr("time.sleep", lambda *a, **kw: None) + + advisory_lock_stderr = ( + "Error: P1002\nThe database server was reached but timed out.\n" + "Context: Timed out trying to acquire a postgres advisory lock " + "(SELECT pg_advisory_lock(72707369)). Elapsed: 10000ms." + ) + + calls = {"n": 0} + + class FakeResult: + stdout = "No pending migrations to apply\n" + stderr = "" + + def fake_run(cmd, *args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise subprocess.CalledProcessError( + returncode=1, + cmd=cmd, + stderr=advisory_lock_stderr, + output="", + ) + return FakeResult() + + monkeypatch.setattr("subprocess.run", fake_run) + + 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" + + +def test_v2_p1002_terminal_raises_with_remediation_hint(monkeypatch, tmp_path): + """v2: 4 consecutive P1002 advisory-lock failures must raise a + RuntimeError whose message names the stale-lock case and includes the + inspection / pg_terminate_backend SQL plus the DIRECT_URL pointer.""" + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setattr( + ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None + ) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setattr("time.sleep", lambda *a, **kw: None) + + advisory_lock_stderr = ( + "Error: P1002\nContext: Timed out trying to acquire a postgres " + "advisory lock (SELECT pg_advisory_lock(72707369)). Elapsed: 10000ms." + ) + + with patch( + "subprocess.run", + side_effect=_fake_migrate_deploy_failure(1, advisory_lock_stderr), + ): + with pytest.raises(RuntimeError) as excinfo: + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + msg = str(excinfo.value) + assert "pg_locks" in msg + assert "pg_terminate_backend" in msg + assert "objid = 72707369" in msg + assert "DIRECT_URL" in msg