diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 493e09e3af1..1468e312ee6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3320,7 +3320,9 @@ def to_dict(self) -> dict: class CommonProxyErrors(str, enum.Enum): db_not_connected_error = ( - "DB not connected. See https://docs.litellm.ai/docs/proxy/virtual_keys" + "DB not connected. This endpoint needs a database; set DATABASE_URL to a " + "PostgreSQL connection string (postgresql://...) to enable it. " + "See https://docs.litellm.ai/docs/proxy/virtual_keys" ) no_llm_router = "No models configured on proxy" not_allowed_access = "Admin-only endpoint. Not allowed to access this." diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 58478db5e2e..89f97aad1d7 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -32,7 +32,7 @@ import os import urllib.parse -from typing import Optional, cast +from typing import Final, Optional, cast from pydantic import AliasChoices, Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -44,6 +44,41 @@ _IAM_ENV_KEY = "IAM_TOKEN_DB_AUTH" _DEFAULT_PG_PORT = "5432" +# schema.prisma pins `provider = "postgresql"`, so these are the only schemes +# Prisma can actually connect with. +SUPPORTED_DB_SCHEMES: Final[frozenset[str]] = frozenset({"postgresql", "postgres"}) +_MISSING_SCHEME = "" + + +def unsupported_db_scheme(database_url: str) -> Optional[str]: + """Return the connection URL scheme when it is not PostgreSQL, else None. + + A `sqlite://` / `mysql://` URL can never connect against the + postgresql-only datasource, but the resulting Prisma failure is opaque and + version-dependent (a confusing migration error, or a startup that never + binds). Callers use this to reject the URL up front with an actionable + error instead. + + A schemeless value (e.g. a malformed DSN like ``user:pass@host/db``) yields + the ``_MISSING_SCHEME`` placeholder rather than the raw URL, so callers that + log the return value never echo embedded credentials. + """ + scheme = urllib.parse.urlsplit(database_url).scheme.lower() + if scheme in SUPPORTED_DB_SCHEMES: + return None + return scheme or _MISSING_SCHEME + + +def unsupported_db_scheme_message(env_var: str, scheme: str) -> str: + """Operator-facing message naming the offending env var and scheme.""" + return ( + f"{env_var} uses unsupported scheme '{scheme}'. LiteLLM's database " + "features (virtual keys, store_model_in_db, spend tracking) require " + "PostgreSQL; use a 'postgresql://' connection string. SQLite and other " + "engines are not supported. " + "See https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + class DatabaseURLSettings(BaseSettings): """Discrete ``DATABASE_*`` env vars, loaded once at process start. @@ -239,6 +274,24 @@ def _password_url( url += f"?schema={schema}" return url + def _raise_for_unsupported_scheme(self) -> None: + """Reject an operator-pinned non-PostgreSQL writer/reader URL. + + The componentized entrypoints (gateway / backend / migrations) call + ``apply_to_env`` and then hand the URL straight to Prisma, bypassing + the CLI's own guard. A pinned URL flows through untouched, so validate + it here too rather than letting Prisma stall on an unusable scheme. + """ + for env_var, url in ( + ("DATABASE_URL", self.database_url), + ("DATABASE_URL_READ_REPLICA", self.database_url_read_replica), + ): + if not url: + continue + bad_scheme = unsupported_db_scheme(url) + if bad_scheme is not None: + raise RuntimeError(unsupported_db_scheme_message(env_var, bad_scheme)) + def apply_to_env(self) -> bool: """Write the assembled URL(s) into ``os.environ``. @@ -246,6 +299,7 @@ def apply_to_env(self) -> bool: password auth that assembled a fresh URL). False means there was nothing to do — an operator-pinned URL, or no discrete fields. """ + self._raise_for_unsupported_scheme() wrote_writer = False writer_url = self.build_writer_url() if writer_url is not None: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 8c3fa952903..d0a043a1821 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1140,6 +1140,25 @@ def run_server( # noqa: PLR0915 os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None ): + from litellm.proxy.db.db_url_settings import ( + unsupported_db_scheme, + unsupported_db_scheme_message, + ) + + for _db_env in ("DATABASE_URL", "DIRECT_URL"): + _candidate_url = os.getenv(_db_env) + if _candidate_url is None: + continue + _bad_scheme = unsupported_db_scheme(_candidate_url) + if _bad_scheme is not None: + print( # noqa + f"\033[1;31mLiteLLM Proxy: " + f"{unsupported_db_scheme_message(_db_env, _bad_scheme)}" + "\033[0m", + file=sys.stderr, + flush=True, + ) + sys.exit(1) try: from litellm.secret_managers.main import get_secret diff --git a/tests/test_litellm/proxy/db/test_db_url_settings.py b/tests/test_litellm/proxy/db/test_db_url_settings.py index b2212068a5b..c8ac39270c0 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -16,7 +16,11 @@ import pytest -from litellm.proxy.db.db_url_settings import DatabaseURLSettings +from litellm.proxy.db.db_url_settings import ( + DatabaseURLSettings, + unsupported_db_scheme, + unsupported_db_scheme_message, +) def _apply() -> bool: @@ -287,3 +291,77 @@ def test_password_reader_uses_own_credentials(monkeypatch): os.environ["DATABASE_URL_READ_REPLICA"] == "postgresql://litellm_ro:ro_pw@reader.example.com:5432/litellm_db" ) + + +@pytest.mark.parametrize( + "url", + [ + "postgresql://u:p@host:5432/db", + "postgres://u:p@host:5432/db", + "POSTGRESQL://u:p@host:5432/db", + "postgresql://host/db?schema=public", + ], +) +def test_unsupported_db_scheme_accepts_postgres(url): + assert unsupported_db_scheme(url) is None + + +@pytest.mark.parametrize( + "url,scheme", + [ + ("sqlite:///data/litellm.db", "sqlite"), + ("sqlite:///./local.db", "sqlite"), + ("mysql://u:p@host:3306/db", "mysql"), + ("mssql://host/db", "mssql"), + ], +) +def test_unsupported_db_scheme_rejects_non_postgres(url, scheme): + assert unsupported_db_scheme(url) == scheme + + +def test_unsupported_db_scheme_does_not_echo_schemeless_credentials(): + """A malformed schemeless DSN must not leak its embedded credentials + through the return value (which callers log).""" + leaky = "litellm:s3cr3t_password@db.internal:5432/litellm" + + result = unsupported_db_scheme(leaky) + + assert result is not None + assert "s3cr3t_password" not in result + assert "db.internal" not in result + + +def test_apply_to_env_rejects_pinned_sqlite_writer(monkeypatch): + """Componentized entrypoints pin DATABASE_URL and call apply_to_env; a + sqlite writer must raise here rather than reach Prisma.""" + monkeypatch.setenv("DATABASE_URL", "sqlite:///data/litellm.db") + + with pytest.raises(RuntimeError, match="sqlite"): + _apply() + + # The bad URL must not have been propagated as a usable connection string. + assert os.environ["DATABASE_URL"] == "sqlite:///data/litellm.db" + + +def test_apply_to_env_rejects_pinned_non_postgres_reader(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") + monkeypatch.setenv( + "DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db" + ) + + with pytest.raises(RuntimeError, match="DATABASE_URL_READ_REPLICA.*mysql"): + _apply() + + +def test_apply_to_env_accepts_pinned_postgres(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@host:5432/db") + + # Operator-pinned URL: nothing reassembled, no error. + assert _apply() is False + + +def test_unsupported_db_scheme_message_names_var_and_scheme(): + msg = unsupported_db_scheme_message("DIRECT_URL", "sqlite") + assert "DIRECT_URL" in msg + assert "sqlite" in msg + assert "postgresql://" in msg diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 34c88e2fd33..2a7c2e1559f 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1483,6 +1483,57 @@ def test_startup_fails_when_db_setup_fails( use_migrate=True, use_v2_resolver=False ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") + def test_startup_exits_on_non_postgres_database_url( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_atexit_register, + mock_subprocess_run, + ): + """A sqlite DATABASE_URL must exit immediately, before any prisma call, + instead of stalling on a migration against the postgresql-only schema.""" + from litellm.proxy.proxy_cli import run_server + + mock_subprocess_run.return_value = MagicMock(returncode=0) + mock_should_update_schema.return_value = True + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env["DATABASE_URL"] = "sqlite:///data/litellm.db" + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + ): + with pytest.raises(SystemExit) as exc_info: + run_server.main( + ["--local", "--skip_server_startup"], standalone_mode=False + ) + assert exc_info.value.code == 1 + mock_setup_database.assert_not_called() + # --- Module-level helpers for worker startup hook tests ---