diff --git a/litellm/repositories/fips_repository.py b/litellm/repositories/fips_repository.py index 23a756e44bf..15c80272483 100644 --- a/litellm/repositories/fips_repository.py +++ b/litellm/repositories/fips_repository.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Final, Protocol, cast # noqa: TID251 # Prisma delegates query_raw and execute_raw dynamically from litellm.proxy.db.create_views import SupportsRawQueries +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry FIPS_MARKER: Final = "fips_mode" _IS_MARKED: Final = 'SELECT 1 FROM "LiteLLM_Config" WHERE param_name = $1' @@ -25,6 +26,10 @@ class _WithWriter(Protocol): @property def writer_db(self) -> object: ... + async def attempt_db_reconnect( + self, *, reason: str, timeout_seconds: float | None = None, lock_timeout_seconds: float | None = None + ) -> bool: ... + @dataclass(frozen=True, slots=True) class FipsDatabaseState: @@ -33,7 +38,8 @@ class FipsDatabaseState: class FipsRepository: - """Reads and writes go to the writer, since a replica may not have the marker yet.""" + """Reads and writes go to the writer, since a replica may not have the marker yet. A read that loses its + connection reconnects and runs once more.""" def __init__(self, prisma_client: _WithWriter) -> None: self._prisma_client: Final = prisma_client @@ -44,10 +50,20 @@ class FipsRepository: return cast(SupportsRawQueries, writer) # cast-ok: Prisma delegates database methods dynamically async def is_marked_fips(self) -> bool: - return bool(await self._writer.query_raw(_IS_MARKED, FIPS_MARKER)) + return await call_with_db_reconnect_retry( + self._prisma_client, self._read_marker, reason="fips_startup_read_marker_failure" + ) async def state(self) -> FipsDatabaseState: """The marker and whether any other LiteLLM row exists, read in one statement so they agree.""" + return await call_with_db_reconnect_retry( + self._prisma_client, self._read_state, reason="fips_startup_read_state_failure" + ) + + async def _read_marker(self) -> bool: + return bool(await self._writer.query_raw(_IS_MARKED, FIPS_MARKER)) + + async def _read_state(self) -> FipsDatabaseState: writer: Final = self._writer tables: Final = tuple(str(row["table_name"]) for row in await writer.query_raw(_OTHER_TABLES)) any_row_in_other_tables: Final = "".join(f" OR EXISTS (SELECT 1 FROM {_quoted(table)})" for table in tables) diff --git a/tests/unit/proxy/common_utils/test_fips.py b/tests/unit/proxy/common_utils/test_fips.py index cccf261fd5c..732cfa55d75 100644 --- a/tests/unit/proxy/common_utils/test_fips.py +++ b/tests/unit/proxy/common_utils/test_fips.py @@ -4,6 +4,7 @@ from contextlib import nullcontext from typing import Final from unittest.mock import Mock +import httpx import pytest from cryptography.exceptions import InternalError from cryptography.hazmat.bindings._rust import openssl as rust_openssl @@ -135,12 +136,16 @@ def test_enforce_fips_mode( class _Database: """Prisma's raw-query edge over a database whose only LiteLLM table besides LiteLLM_Config is the key table.""" - def __init__(self, *, marked_fips: bool, holds_rows: bool, readable: bool = True) -> None: + def __init__(self, *, marked_fips: bool, holds_rows: bool, readable: bool = True, dropped_reads: int = 0) -> None: self.marked_fips = marked_fips self.holds_rows = holds_rows self.readable = readable + self.dropped_reads = dropped_reads async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + if self.dropped_reads: + self.dropped_reads -= 1 + raise httpx.ReadError("connection dropped") if not self.readable: raise RuntimeError('relation "LiteLLM_Config" does not exist') if "information_schema.tables" in query: @@ -158,6 +163,11 @@ class _Prisma: def __init__(self, database: _Database) -> None: self.writer_db: Final = database + async def attempt_db_reconnect( + self, *, reason: str, timeout_seconds: float | None = None, lock_timeout_seconds: float | None = None + ) -> bool: + return True + @pytest.mark.asyncio @pytest.mark.parametrize( @@ -170,6 +180,19 @@ class _Prisma: (False, _Database(marked_fips=True, holds_rows=True), "set up with general_settings.fips_mode: true", True), (False, _Database(marked_fips=False, holds_rows=True), None, False), (False, _Database(marked_fips=True, holds_rows=True, readable=False), None, True), + ( + True, + _Database(marked_fips=False, holds_rows=True, dropped_reads=1), + "already holds data written without it", + False, + ), + ( + False, + _Database(marked_fips=True, holds_rows=True, dropped_reads=1), + "set up with general_settings.fips_mode: true", + True, + ), + (False, _Database(marked_fips=True, holds_rows=True, dropped_reads=2), None, True), ], ids=[ "fips-on-new-database", @@ -179,6 +202,9 @@ class _Prisma: "fips-off-fips-database", "fips-off-used-database", "fips-off-unreadable-database", + "fips-on-used-database-after-a-dropped-connection", + "fips-off-fips-database-after-a-dropped-connection", + "fips-off-fips-database-dropped-twice", ], ) async def test_enforce_fips_database(