From 2b8b9502d91af5fa5247f13c0b63b669bb54b329 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 21 Apr 2026 15:53:07 -0700 Subject: [PATCH] [Fix] v2 resolver: swallow non-connection DB errors; wrap resolve failures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses two further Greptile findings: - `_warn_if_db_ahead_of_head` only caught `psycopg.OperationalError`. Non-connection DB errors (e.g. `InsufficientPrivilege` / 42501 if the runtime DB user lacks SELECT on `_prisma_migrations`) would propagate uncaught and crash startup — contradicting the docstring's "informational only, never blocks" guarantee. Widen the catch to `psycopg.DatabaseError` so all DB-layer errors are swallowed. - In the P3009 and P3018 idempotent-recovery paths, the call to `_resolve_specific_migration(name)` was not wrapped in its own try/except. Being inside an active `except CalledProcessError` handler, a new `CalledProcessError` from the resolve call would NOT re-enter the same handler — it would propagate out as `CalledProcessError`, past `proxy_cli.py`'s `except RuntimeError`, crashing startup with an unhandled traceback instead of the intended clean `sys.exit(2)`. Wrap both call sites to convert to RuntimeError. Adds unit tests for both behaviors. --- .../litellm_proxy_extras/utils.py | 41 +++++++++-- .../tests/test_setup_database_fail_fast.py | 72 +++++++++++++++++++ 2 files changed, 108 insertions(+), 5 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index a234ad18ba..369b656193 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -470,7 +470,11 @@ class ProxyExtrasDBManager: ).fetchall() except psycopg.errors.UndefinedTable: return - except psycopg.OperationalError: + except (psycopg.OperationalError, psycopg.DatabaseError): + # Swallow connection failures AND any other DB-layer error + # (e.g. InsufficientPrivilege if the runtime user lacks SELECT + # on _prisma_migrations). This is an informational check — + # never block startup on it. return applied = {r[0] for r in rows} @@ -589,8 +593,24 @@ class ProxyExtrasDBManager: subprocess.CalledProcessError, subprocess.TimeoutExpired, ): - pass - ProxyExtrasDBManager._resolve_specific_migration(name) + pass # may already be rolled-back + try: + ProxyExtrasDBManager._resolve_specific_migration(name) + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ) as resolve_err: + # We're already inside the outer + # `except CalledProcessError` handler — + # re-raising CalledProcessError from here + # would escape as itself, bypassing + # proxy_cli.py's `except RuntimeError`. + raise RuntimeError( + f"Failed to mark migration {name} as applied " + f"after idempotent recovery. Manual " + f"intervention may be required.\n\n" + f"Detail: {resolve_err}" + ) from resolve_err continue raise RuntimeError( "Database migration failed and cannot be auto-recovered. " @@ -622,8 +642,19 @@ class ProxyExtrasDBManager: subprocess.CalledProcessError, subprocess.TimeoutExpired, ): - pass - ProxyExtrasDBManager._resolve_specific_migration(name) + pass # may already be rolled-back + try: + ProxyExtrasDBManager._resolve_specific_migration(name) + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ) as resolve_err: + raise RuntimeError( + f"Failed to mark migration {name} as applied " + f"after idempotent recovery. Manual " + f"intervention may be required.\n\n" + f"Detail: {resolve_err}" + ) from resolve_err continue raise RuntimeError( 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 573137c90e..8d66bf872d 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -144,6 +144,78 @@ def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_pat ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True) +def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path): + """_warn_if_db_ahead_of_head must never raise — it's informational. + + Non-connection DB errors (e.g. InsufficientPrivilege from a user + without SELECT on _prisma_migrations) must be caught, not propagated. + """ + import psycopg + + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + + class _FakeConn: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def execute(self, *a, **kw): + # Simulate an InsufficientPrivilege (subclass of DatabaseError). + raise psycopg.errors.InsufficientPrivilege("permission denied") + + def _fake_connect(*a, **kw): + return _FakeConn() + + monkeypatch.setattr("psycopg.connect", _fake_connect) + + # Must not raise. + ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path)) + + +def test_v2_resolve_specific_migration_failure_raises_runtime_error( + monkeypatch, tmp_path +): + """If marking a migration as applied fails inside P3009 idempotent + recovery, the subprocess error must be re-raised as RuntimeError so + proxy_cli.py catches it cleanly (instead of leaking CalledProcessError).""" + 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( + ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None + ) + + # First call: migrate deploy -> P3009 idempotent error. + # Recovery path tries _resolve_specific_migration; that also raises. + def _failing_resolve(*a, **kw): + raise subprocess.CalledProcessError( + returncode=1, + cmd="prisma migrate resolve --applied", + stderr="resolve failed", + output="", + ) + + monkeypatch.setattr( + ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve + ) + + stderr = ( + "Error: P3009\nMigration `20260101000000_some_migration` failed\n" + "relation already exists" + ) + with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)): + with pytest.raises( + RuntimeError, match="Failed to mark migration .* as applied" + ): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path): """v2 must never call _resolve_all_migrations — that's the bug it fixes.""" monkeypatch.setattr(