From 99f99cfb46a46eaa65333c44304e12e448a21bbb Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 19 Sep 2026 18:10:52 -0700 Subject: [PATCH] fix(proxy): stop the boot when the requested master key migration fails, unless allow_requests_on_db_unavailable tolerates the outage --- litellm/proxy/db/master_key_migration.py | 15 +++-- litellm/proxy/proxy_server.py | 1 + .../proxy/auth/test_master_key_boot_check.py | 34 +++++++++-- .../proxy/db/test_master_key_migration.py | 59 ++++++++++++++++--- tests/test_litellm/proxy/test_proxy_server.py | 30 ++++++++++ 5 files changed, 123 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/db/master_key_migration.py b/litellm/proxy/db/master_key_migration.py index 78a00cf9ce9..d100554201a 100644 --- a/litellm/proxy/db/master_key_migration.py +++ b/litellm/proxy/db/master_key_migration.py @@ -174,7 +174,7 @@ class Migrated: @dataclass(frozen=True, slots=True) class MigrationFailed: - error: str + error: Exception MigrationOutcome = NothingToMigrate | Migrated | MigrationFailed @@ -186,17 +186,21 @@ async def migrate_if_requested( master_key: str | None, connected_database: Callable[[], SupportsRawQueries | None], log: Callable[[str], None], + raise_unless_tolerated: Callable[[Exception], None], ) -> MigrationOutcome | None: previous_master_key: Final = environ.get(MIGRATE_FROM_MASTER_KEY_ENV_VAR) if previous_master_key is None or master_key is None: return None - return await migrate_from_previous_master_key( + outcome: Final = await migrate_from_previous_master_key( previous_master_key=previous_master_key, master_key=master_key, salt_key_is_set=SALT_KEY_ENV_VAR in environ, database=connected_database(), log=log, ) + if isinstance(outcome, MigrationFailed): + raise_unless_tolerated(outcome.error) + return outcome async def migrate_from_previous_master_key( @@ -234,8 +238,8 @@ async def _migrate_or_failure( database=database, log=log, ) - except Exception as error: # noqa: BLE001 # the proxy tolerates a database outage at boot, so the migration must too - return MigrationFailed(error=f"{type(error).__name__}: {error}"[:300]) + except Exception as error: # noqa: BLE001 # a value, so the boot applies its own database outage rule to it + return MigrationFailed(error=error) async def _migrate( @@ -293,8 +297,9 @@ def describe_outcome(outcome: MigrationOutcome) -> str: "the proxy to migrate them." ) case MigrationFailed(error=error): + cause: Final = f"{type(error).__name__}: {error}"[:300] return ( - f"Could not migrate stored values from the {MIGRATE_FROM_MASTER_KEY_ENV_VAR} key ({error}). Values " + f"Could not migrate stored values from the {MIGRATE_FROM_MASTER_KEY_ENV_VAR} key ({cause}). Values " "still encrypted with the previous key cannot be read until the migration succeeds. Keep " f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} set and restart the proxy once the database is reachable." ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0b15ea902ac..c0c8dfeaeee 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1267,6 +1267,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: master_key=master_key, connected_database=lambda: None if prisma_client is None else prisma_client.writer_db, log=verbose_proxy_logger.warning, + raise_unless_tolerated=PrismaDBExceptionHandler.handle_db_exception, ) if prisma_client is not None: diff --git a/tests/test_litellm/proxy/auth/test_master_key_boot_check.py b/tests/test_litellm/proxy/auth/test_master_key_boot_check.py index bf208bf235b..0b55d045de5 100644 --- a/tests/test_litellm/proxy/auth/test_master_key_boot_check.py +++ b/tests/test_litellm/proxy/auth/test_master_key_boot_check.py @@ -254,10 +254,36 @@ def test_refusal_never_tells_a_user_with_an_exported_key_to_append_to_the_env_fi assert "wins over .env" in text -def test_unset_key_refusal_says_nothing_supplied_one(): - text = render_refusal(_refusal(reason=UnsafeMasterKeyReason.NOT_SET, source=EnvironmentSource())) - - assert "Neither general_settings.master_key nor" in text +@pytest.mark.parametrize( + ("reason", "source", "source_line"), + [ + ( + UnsafeMasterKeyReason.NOT_SET, + EnvironmentSource(), + f"Neither general_settings.master_key nor the {MASTER_KEY_ENV_VAR} environment variable is set.", + ), + ( + UnsafeMasterKeyReason.PUBLICLY_KNOWN, + EnvironmentSource(), + f"It comes from the {MASTER_KEY_ENV_VAR} environment variable.", + ), + ( + UnsafeMasterKeyReason.NOT_SET, + ConfigFileSource(config_file_path="/app/config.yaml"), + "general_settings.master_key in /app/config.yaml is blank, " + "or points at an environment variable that is not set.", + ), + ( + UnsafeMasterKeyReason.EMPTY, + ConfigFileSource(config_file_path="/app/config.yaml"), + "It comes from general_settings.master_key in /app/config.yaml.", + ), + ], +) +def test_refusal_says_where_the_unsafe_key_came_from( + reason: UnsafeMasterKeyReason, source: ConfigFileSource | EnvironmentSource, source_line: str +): + assert render_refusal(_refusal(reason=reason, source=source)).splitlines()[1] == source_line def test_migration_steps_appear_only_when_the_database_needs_them(): diff --git a/tests/test_litellm/proxy/db/test_master_key_migration.py b/tests/test_litellm/proxy/db/test_master_key_migration.py index 4cc9a652969..9c0fc163b9f 100644 --- a/tests/test_litellm/proxy/db/test_master_key_migration.py +++ b/tests/test_litellm/proxy/db/test_master_key_migration.py @@ -439,21 +439,61 @@ async def test_encrypted_empty_string_is_migrated_like_any_other_value(): assert decrypt_if_encrypted_with(str(tables["LiteLLM_MCPUserCredentials"][0]["credential_b64"]), NEW_KEY) == "" -@pytest.mark.asyncio -async def test_database_error_during_the_migration_is_reported_instead_of_crashing_the_boot(): - class _DatabaseIsDown(_FakeDatabase): - async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: - raise ConnectionError("Can't reach database server") +class _DatabaseIsDown(_FakeDatabase): + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + raise ConnectionError("Can't reach database server") + +@pytest.mark.asyncio +async def test_database_error_during_the_migration_comes_back_as_a_value_and_is_logged(): outcome, logged = await _run(_DatabaseIsDown(_seeded_tables())) - assert outcome == MigrationFailed(error="ConnectionError: Can't reach database server") + assert isinstance(outcome, MigrationFailed) + assert isinstance(outcome.error, ConnectionError) assert len(logged) == 1 assert "ConnectionError: Can't reach database server" in logged[0] assert f"Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set" in logged[0] assert "ou may now delete" not in logged[0] +def _raise(error: Exception) -> None: + raise error + + +def _tolerate(error: Exception) -> None: + return None + + +@pytest.mark.asyncio +async def test_boot_stops_on_a_failed_migration_when_the_outage_is_not_tolerated(): + logged: list[str] = [] + + with pytest.raises(ConnectionError, match="Can't reach database server"): + await migrate_if_requested( + environ={MIGRATE_FROM_MASTER_KEY_ENV_VAR: PREVIOUS_KEY}, + master_key=NEW_KEY, + connected_database=lambda: _DatabaseIsDown(_seeded_tables()), + log=logged.append, + raise_unless_tolerated=_raise, + ) + + assert len(logged) == 1 + assert f"Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set" in logged[0] + + +@pytest.mark.asyncio +async def test_boot_continues_past_a_failed_migration_when_the_outage_is_tolerated(): + outcome = await migrate_if_requested( + environ={MIGRATE_FROM_MASTER_KEY_ENV_VAR: PREVIOUS_KEY}, + master_key=NEW_KEY, + connected_database=lambda: _DatabaseIsDown(_seeded_tables()), + log=lambda line: None, + raise_unless_tolerated=_tolerate, + ) + + assert isinstance(outcome, MigrationFailed) + + @pytest.mark.asyncio @pytest.mark.parametrize("previous_master_key", [PREVIOUS_KEY, ""]) async def test_boot_migrates_from_the_environment_variable_to_the_running_master_key(previous_master_key: str): @@ -476,6 +516,7 @@ async def test_boot_migrates_from_the_environment_variable_to_the_running_master master_key=NEW_KEY, connected_database=lambda: _FakeDatabase(tables), log=logged.append, + raise_unless_tolerated=_raise, ) assert outcome == Migrated(migrated=1, remaining=0) @@ -510,7 +551,11 @@ async def test_boot_leaves_the_database_alone_unless_a_migration_was_requested_a return _DatabaseThatMustNotBeTouched() result = await migrate_if_requested( - environ=environ, master_key=master_key, connected_database=connected_database, log=logged.append + environ=environ, + master_key=master_key, + connected_database=connected_database, + log=logged.append, + raise_unless_tolerated=_raise, ) assert result is outcome diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 191fa862f6c..32f18ce1e4e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1741,6 +1741,36 @@ async def test_proxy_startup_says_a_lingering_migrate_from_variable_can_be_delet assert "you may now delete LITELLM_MIGRATE_FROM_MASTER_KEY" in notices[0] +class _PrismaClientWhoseDatabaseRejectsQueries: + class _Database: + async def query_raw(self, query, *args): + raise RuntimeError("permission denied for table LiteLLM_CredentialsTable") + + writer_db = _Database() + + +@pytest.mark.asyncio +async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatch, tmp_path, caplog): + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries()) + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "sk-1234") + + with ( + caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"), + pytest.raises(RuntimeError, match="permission denied"), + ): + async with proxy_startup_event(FastAPI()): + pass + + notices = [record.getMessage() for record in caplog.records if "LITELLM_MIGRATE_FROM_MASTER_KEY" in record.message] + assert len(notices) == 1 + assert "Could not migrate stored values" in notices[0] + + @pytest.mark.asyncio async def test_proxy_startup_names_the_config_file_that_set_the_unsafe_key(monkeypatch, tmp_path): from fastapi import FastAPI