fix(proxy-extras): name the database error when the Lens rename check cannot run (#45341)

* fix(proxy-extras): name the database error when the Lens rename check cannot run

The Lens rename check runs before prisma db push and raised a generic
"Cannot verify Lens data safety" RuntimeError on any psycopg error, with
the cause only in the chained exception. proxy_cli prints the message and
exits, so a Postgres at its connection limit read as a Lens problem.

The raised message now ends with the redacted database error text, and
the connection opener is an injected parameter so the unit test drives
the check without a database.

* test(proxy-extras): drop the Lens check test class docstring

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -11,7 +11,7 @@ import time
 from collections.abc import Callable
 from dataclasses import dataclass, replace
 from pathlib import Path
-from typing import TYPE_CHECKING, Final, Optional, Union
+from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
 from urllib.parse import unquote, urlsplit
 
 from litellm_proxy_extras import prisma_toolchain
@@ -35,6 +35,12 @@ if TYPE_CHECKING:
     import psycopg.sql
 
 
+class LensCheckConnect(Protocol):
+    def __call__(
+        self, conninfo: str, *, connect_timeout: int, autocommit: bool
+    ) -> "psycopg.Connection[tuple[object, ...]]": ...
+
+
 def str_to_bool(value: Optional[str]) -> bool:
     if value is None:
         return False
@@ -243,11 +249,9 @@ def _redact_credentials(text: str) -> str:
     passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
     alternation: Final = "|".join(re.escape(password) for password in passwords)
     password_pattern: Final = (
-        re.compile(rf"(?P<lead>:|password=)(?:{alternation})(?=@|&|$|[\s'\"\]),])", re.IGNORECASE)
-        if passwords
-        else None
+        re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
     )
-    result: Final = password_pattern.sub(rf"\g<lead>{_REDACTED}", text) if password_pattern is not None else text
+    result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
     return _secret_shape_redactor()(result)
 
 
@@ -690,7 +694,7 @@ class ProxyExtrasDBManager:
                     )
 
     @staticmethod
-    def raise_if_lens_rename_pending() -> None:
+    def raise_if_lens_rename_pending(connect: LensCheckConnect | None = None) -> None:
         database_url: Final = os.environ.get("DATABASE_URL")
         if not database_url:
             return
@@ -698,8 +702,9 @@ class ProxyExtrasDBManager:
             import psycopg
         except ImportError as exc:
             raise RuntimeError("Install psycopg to verify Lens data safety before prisma db push.") from exc
+        open_connection: Final = connect if connect is not None else psycopg.connect
         try:
-            with psycopg.connect(
+            with open_connection(
                 ProxyExtrasDBManager._strip_prisma_query_params(database_url), connect_timeout=10, autocommit=True
             ) as connection:
                 legacy: Final = connection.execute(
@@ -710,7 +715,8 @@ class ProxyExtrasDBManager:
                 ).fetchone()
         except psycopg.Error as exc:
             raise RuntimeError(
-                "Cannot verify Lens data safety; refusing prisma db push. Check database connectivity and psycopg installation."
+                "Cannot verify Lens data safety; refusing prisma db push. "
+                f"The database check failed: {_redact_credentials(str(exc)).strip()}"
             ) from exc
         if legacy is not None:
             raise RuntimeError(

tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py

@@ -655,6 +655,49 @@ class _FakePsycopgConn:
         return _FakeCursor()
 
 
+class TestLensRenamePendingCheck:
+    _DATABASE_URL: Final = "postgresql://litellm:hunter2@localhost:5432/litellm"
+
+    def test_database_failure_text_reaches_the_raised_message(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        import psycopg
+
+        monkeypatch.setenv("DATABASE_URL", self._DATABASE_URL)
+
+        def refuse(conninfo: str, *, connect_timeout: int, autocommit: bool) -> NoReturn:
+            raise psycopg.OperationalError("FATAL:  sorry, too many clients already")
+
+        with pytest.raises(RuntimeError) as err:
+            ProxyExtrasDBManager.raise_if_lens_rename_pending(connect=refuse)
+        assert "FATAL:  sorry, too many clients already" in str(err.value)
+
+    def test_password_libpq_echoes_is_redacted_from_the_message(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        import psycopg
+
+        password: Final = "p%zzword"
+        database_url: Final = f"postgresql://litellm:{password}@localhost:5432/litellm"
+        monkeypatch.setenv("DATABASE_URL", database_url)
+        monkeypatch.delenv("DIRECT_URL", raising=False)
+        with pytest.raises(psycopg.Error) as libpq:
+            psycopg.connect(database_url, connect_timeout=10, autocommit=True)
+
+        with pytest.raises(RuntimeError) as err:
+            ProxyExtrasDBManager.raise_if_lens_rename_pending()
+        assert password not in str(err.value)
+        assert str(err.value).endswith(str(libpq.value).strip().replace(password, "REDACTED"))
+
+    def test_legacy_tables_keep_their_own_message(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("DATABASE_URL", self._DATABASE_URL)
+        executed: Final[list[tuple[str, tuple[str]]]] = []
+
+        def legacy_tables_present(conninfo: str, *, connect_timeout: int, autocommit: bool) -> _FakePsycopgConn:
+            return _FakePsycopgConn(executed)
+
+        with pytest.raises(RuntimeError) as err:
+            ProxyExtrasDBManager.raise_if_lens_rename_pending(connect=legacy_tables_present)
+        assert str(err.value).startswith("Legacy Lens tables exist.")
+        assert executed[0][1] == ("public",)
+
+
 class TestSpendLogsPartitionDetectionSchemaScope:
     """A same-named LiteLLM_SpendLogs in another schema must not trip the
     detector: the catalog lookup has to be scoped to Prisma's target schema."""