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

@@ -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."
             )

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:

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():

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

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