diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 7f8db72a1a6..939d97e3ec2 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -300,6 +300,19 @@ def connection_params_from_url(url: str) -> Mapping[str, str | int | float]: ) +def schema_from_database_url(url: str) -> str | None: + """The Postgres schema Prisma resolves queries against for ``url``. + + A pre-existing ``DATABASE_URL`` an operator pins directly (rather than one + assembled from discrete ``DATABASE_*`` vars) carries its ``?schema=`` as + part of the URL itself, never mirrored into ``DATABASE_SCHEMA``. Anything + that needs the schema Prisma is actually querying must read it from here + too, not assume ``DATABASE_SCHEMA`` is set. + """ + pairs: Final = urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query) + return next((value for key, value in pairs if key == "schema"), None) + + # A re-minted token URL replaces a URL to the same database, so unlike the # reader allowlist it may also carry ``options``: that is where the server-side # timeouts (statement, lock, idle-in-transaction) live, and a refresh that diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 267c87694b9..1534992949b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -182,6 +182,7 @@ from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.db_url_settings import ( DatabaseURLSettings, add_missing_query_params, + schema_from_database_url, token_refresh_params_from_url, ) from litellm.proxy.db.exception_handler import ( @@ -4931,7 +4932,9 @@ class PrismaClient: ] required_view: Final = "LiteLLM_VerificationTokenView" expected_views_str: Final = ", ".join(f"'{view}'" for view in expected_views) - pg_schema: Final = os.getenv("DATABASE_SCHEMA", "public") + pg_schema: Final = ( + os.getenv("DATABASE_SCHEMA") or schema_from_database_url(os.getenv("DATABASE_URL", "")) or "public" + ) ret: Final[Sequence[_ViewCountRow]] = await self.db.query_raw(f""" WITH existing_views AS ( SELECT viewname diff --git a/tests/unit/proxy/db/test_db_url_settings.py b/tests/unit/proxy/db/test_db_url_settings.py index 99957447248..57b48d94a9f 100644 --- a/tests/unit/proxy/db/test_db_url_settings.py +++ b/tests/unit/proxy/db/test_db_url_settings.py @@ -34,6 +34,7 @@ from pydantic import ValidationError from litellm.proxy.db.db_url_settings import ( PG_SSL_REQUEST, DatabaseURLSettings, + schema_from_database_url, token_refresh_params_from_url, translate_libpq_ssl_params, unsupported_db_scheme, @@ -964,6 +965,18 @@ def test_token_refresh_params_keep_the_options_carrying_the_server_timeouts() -> } +def test_schema_from_database_url_reads_a_pinned_schema_query_param(): + assert schema_from_database_url("postgresql://db.example.com:5432/litellm_db?schema=litellm") == "litellm" + + +def test_schema_from_database_url_returns_none_without_a_schema_param(): + assert schema_from_database_url("postgresql://db.example.com:5432/litellm_db?sslmode=require") is None + + +def test_schema_from_database_url_returns_none_for_an_empty_url(): + assert schema_from_database_url("") is None + + def _issue_cert( subject: str, issuer: x509.Certificate | None, issuer_key: ec.EllipticCurvePrivateKey | None, ca: bool ) -> tuple[x509.Certificate, ec.EllipticCurvePrivateKey]: diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index 05c4f9d8a67..da1c2e4fc04 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -427,6 +427,38 @@ async def test_check_view_exists_raises_when_query_raw_fails( await prisma_client.check_view_exists() +@pytest.mark.asyncio +async def test_check_view_exists_checks_the_schema_pinned_on_database_url( + prisma_client: PrismaClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A DATABASE_URL pinned directly (no DATABASE_SCHEMA env var) must still + be checked against its own ``?schema=`` rather than defaulting to public.""" + monkeypatch.delenv("DATABASE_SCHEMA", raising=False) + monkeypatch.setenv( + "DATABASE_URL", "postgresql://db.example.com:5432/app?schema=litellm" + ) + prisma_client.db.query_raw = AsyncMock( + return_value=[{"view_count": 8, "view_names": ["LiteLLM_VerificationTokenView"]}] + ) + await prisma_client.check_view_exists() + assert "schemaname = 'litellm'" in prisma_client.db.query_raw.await_args.args[0] + + +@pytest.mark.asyncio +async def test_check_view_exists_falls_back_to_public_without_a_pinned_schema( + prisma_client: PrismaClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("DATABASE_SCHEMA", raising=False) + monkeypatch.setenv("DATABASE_URL", "postgresql://db.example.com:5432/app") + prisma_client.db.query_raw = AsyncMock( + return_value=[{"view_count": 8, "view_names": ["LiteLLM_VerificationTokenView"]}] + ) + await prisma_client.check_view_exists() + assert "schemaname = 'public'" in prisma_client.db.query_raw.await_args.args[0] + + @pytest.mark.asyncio async def test_get_data_token_find_unique_returns_record( prisma_client: PrismaClient,