diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index 141ce92f172..5ea9cba8018 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -1,15 +1,50 @@ -from typing import Any, Final +from typing import Any, Final, Protocol from litellm import verbose_logger _db = Any + +class SupportsExecuteRaw(Protocol): + """The one database operation create_view_tolerating_race needs. + + Narrower than the `_db = Any` the rest of this module still uses, so the + helper's contract is checkable at its call sites without retyping every + function here. + """ + + async def execute_raw(self, query: str, *args: object) -> int: ... + + # Markers that indicate a view/relation does not yet exist in the database. # Keeping these in one place avoids repeating the check across all view blocks # and prevents overly broad matches (e.g. bare 'undefined' would also match # 'undefined function' or 'column undefined_col referenced in query'). _VIEW_NOT_FOUND_MARKERS: Final = ("does not exist", "no such table", "undefined table") +# Markers for the inverse condition: another replica created the view between +# our existence probe and our CREATE. +_VIEW_ALREADY_EXISTS_MARKERS: Final = ("already exists", "duplicate object", "duplicate table") + + +async def create_view_tolerating_race(db: SupportsExecuteRaw, view_name: str, ddl: str) -> None: + """ + Create a view, treating "a concurrent creator won" as success. + + Every replica booting against the same fresh database observes the view as + absent and issues the CREATE; Postgres fails all but one with a + duplicate-object error. The desired end state is still reached, so losing + that race is success. Without this, the loser's exception propagates out of + a detached startup task and the remaining views are never created. + """ + try: + await db.execute_raw(ddl) + verbose_logger.debug("%s Created!", view_name) + except Exception as e: + if not any(marker in str(e).lower() for marker in _VIEW_ALREADY_EXISTS_MARKERS): + raise + verbose_logger.debug("%s already created by a concurrent replica", view_name) + async def create_missing_views(db: _db): """ @@ -34,7 +69,10 @@ async def create_missing_views(db: _db): if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS): raise # If an error occurs, the view does not exist, so create it - await db.execute_raw(""" + await create_view_tolerating_race( + db, + "LiteLLM_VerificationTokenView", + """ CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -46,9 +84,8 @@ async def create_missing_views(db: _db): FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id LEFT JOIN "LiteLLM_ProjectTable" p ON v.project_id = p.project_id; - """) - - verbose_logger.debug("LiteLLM_VerificationTokenView Created!") + """, + ) try: await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpend" LIMIT 1""") @@ -69,9 +106,7 @@ async def create_missing_views(db: _db): GROUP BY DATE("startTime"); """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("MonthlyGlobalSpend Created!") + await create_view_tolerating_race(db, "MonthlyGlobalSpend", sql_query) try: await db.query_raw("""SELECT 1 FROM "Last30dKeysBySpend" LIMIT 1""") @@ -100,9 +135,7 @@ async def create_missing_views(db: _db): ORDER BY total_spend DESC; """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("Last30dKeysBySpend Created!") + await create_view_tolerating_race(db, "Last30dKeysBySpend", sql_query) try: await db.query_raw("""SELECT 1 FROM "Last30dModelsBySpend" LIMIT 1""") @@ -126,9 +159,7 @@ async def create_missing_views(db: _db): ORDER BY total_spend DESC; """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("Last30dModelsBySpend Created!") + await create_view_tolerating_race(db, "Last30dModelsBySpend", sql_query) try: await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerKey" LIMIT 1""") verbose_logger.debug("MonthlyGlobalSpendPerKey Exists!") @@ -150,9 +181,7 @@ async def create_missing_views(db: _db): DATE("startTime"), api_key; """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("MonthlyGlobalSpendPerKey Created!") + await create_view_tolerating_race(db, "MonthlyGlobalSpendPerKey", sql_query) try: await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1""") verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Exists!") @@ -176,9 +205,7 @@ async def create_missing_views(db: _db): "user", api_key; """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Created!") + await create_view_tolerating_race(db, "MonthlyGlobalSpendPerUserPerKey", sql_query) try: await db.query_raw("""SELECT 1 FROM "DailyTagSpend" LIMIT 1""") @@ -197,9 +224,7 @@ async def create_missing_views(db: _db): FROM "LiteLLM_SpendLogs" s GROUP BY individual_request_tag, DATE(s."startTime"); """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("DailyTagSpend Created!") + await create_view_tolerating_race(db, "DailyTagSpend", sql_query) try: await db.query_raw("""SELECT 1 FROM "Last30dTopEndUsersSpend" LIMIT 1""") @@ -218,9 +243,7 @@ async def create_missing_views(db: _db): ORDER BY total_spend DESC LIMIT 100; """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("Last30dTopEndUsersSpend Created!") + await create_view_tolerating_race(db, "Last30dTopEndUsersSpend", sql_query) async def should_create_missing_views(db: _db) -> bool: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8a1fae42789..ef0376ade84 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -106,6 +106,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_c from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( create_missing_views, + create_view_tolerating_race, should_create_missing_views, ) from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter @@ -3273,7 +3274,10 @@ class PrismaClient: ## check if required view exists ## if ret[0]["view_names"] and required_view not in ret[0]["view_names"]: await self.health_check() # make sure we can connect to db - await self.db.execute_raw(""" + await create_view_tolerating_race( + self.db, + "LiteLLM_VerificationTokenView", + """ CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -3283,9 +3287,8 @@ class PrismaClient: t.rpm_limit AS team_rpm_limit FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; - """) - - verbose_proxy_logger.info("LiteLLM_VerificationTokenView Created in DB!") + """, + ) else: should_create_views: Final = await should_create_missing_views(db=self.db) if should_create_views: diff --git a/tests/test_litellm/proxy/db/test_create_views.py b/tests/test_litellm/proxy/db/test_create_views.py index c0c09d0137b..ecc6d70123e 100644 --- a/tests/test_litellm/proxy/db/test_create_views.py +++ b/tests/test_litellm/proxy/db/test_create_views.py @@ -189,3 +189,66 @@ async def test_create_views_creates_view_on_undefined_table_error(): await create_missing_views(mock_db) mock_db.execute_raw.assert_called_once() + + +# Every view create_missing_views is responsible for. Hard-coded rather than +# derived from the module, so adding a view without guarding it fails here. +EXPECTED_VIEW_COUNT = 8 + + +@pytest.mark.asyncio +async def test_create_views_tolerates_a_concurrent_creator_on_every_view(): + """A replica that loses the CREATE race must attempt every view regardless. + + Regression: two proxy pods booting on a fresh DB both see every view as + absent and both issue the CREATE, and Postgres fails the loser with a + duplicate-object error on whichever views the winner got to first. Any + creation site still calling execute_raw unguarded re-raises that error and + aborts the rest of the function. + + Every CREATE loses here, which is what pins the guard to all of them: an + earlier version of this fix converted only the first and the last site and + still died on MonthlyGlobalSpend against a real Postgres. Counting the + attempts is the assertion, because a partial fix simply stops early. + """ + from litellm.proxy.db.create_views import create_missing_views + + mock_db = MagicMock() + mock_db.query_raw = AsyncMock(side_effect=Exception("relation does not exist")) + mock_db.execute_raw = AsyncMock( + side_effect=Exception('relation "some_view" already exists') + ) + + await create_missing_views(mock_db) + + assert mock_db.execute_raw.await_count == EXPECTED_VIEW_COUNT, ( + f"every view must still be attempted when the replica loses every race; " + f"got {mock_db.execute_raw.await_count} of {EXPECTED_VIEW_COUNT}, so a " + f"creation site is still unguarded and aborted the rest" + ) + + +@pytest.mark.asyncio +async def test_create_views_reraises_genuine_ddl_error(): + """An already-exists guard must not swallow real DDL failures.""" + from litellm.proxy.db.create_views import create_missing_views + + mock_db = MagicMock() + mock_db.query_raw = AsyncMock(side_effect=Exception("relation does not exist")) + mock_db.execute_raw = AsyncMock(side_effect=Exception("syntax error at or near")) + + with pytest.raises(Exception, match="syntax error"): + await create_missing_views(mock_db) + + +@pytest.mark.asyncio +async def test_create_view_tolerating_race_swallows_only_already_exists(): + from litellm.proxy.db.create_views import create_view_tolerating_race + + mock_db = MagicMock() + mock_db.execute_raw = AsyncMock(side_effect=Exception("duplicate object")) + await create_view_tolerating_race(mock_db, "SomeView", "CREATE VIEW ...") + + mock_db.execute_raw = AsyncMock(side_effect=Exception("permission denied")) + with pytest.raises(Exception, match="permission denied"): + await create_view_tolerating_race(mock_db, "SomeView", "CREATE VIEW ...")