diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 1fd292b8137..72462fa9ae4 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/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:|password=)(?:{alternation})(?=@|&|$|[\s'\"\]),])", re.IGNORECASE) - if passwords - else None + re.compile(rf"(?{_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( diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index 70c6ec44790..7edb531a63b 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/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."""