From 1d50d1ad3b9caf69f05d7c2209ec2862da49171e Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 17 Sep 2026 23:40:59 +0000 Subject: [PATCH] fix(proxy): rewrite every model allowlist in one statement on rename Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_allowlist_rename_sync.py | 73 ++++++++------- .../test_model_management_endpoints.py | 88 +++++++++++-------- 2 files changed, 89 insertions(+), 72 deletions(-) diff --git a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py index 1e92c383f01..f93312f7a37 100644 --- a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py +++ b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py @@ -8,37 +8,36 @@ alone denies the new name while the old entry grants a name nothing serves any m from collections.abc import Callable from dataclasses import dataclass +from types import MappingProxyType from typing import Final from pydantic import BaseModel from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.management_helpers.access_group_model_sync import RawExecutor, raw_executor, still_backed +from litellm.proxy.management_helpers.access_group_model_sync import raw_executor, still_backed from litellm.router import Router class _TouchedRow(BaseModel): + kind: str object_id: str team_alias: str | None = None @dataclass(frozen=True, slots=True) class _AllowlistTable: + kind: str table: str - returning: str + id_column: str cache_keys: Callable[[_TouchedRow], tuple[str, ...]] + alias_column: str | None = None - def replace_sql(self) -> str: + def update_cte(self, set_clause: str, where_clause: str) -> str: + alias: Final = f'"{self.alias_column}"' if self.alias_column else "NULL::text" return ( - f'UPDATE "{self.table}" SET "models" = array_replace(array_remove("models", $2), $1, $2) ' - f'WHERE $1 = ANY("models") RETURNING {self.returning}' - ) - - def append_sql(self) -> str: - return ( - f'UPDATE "{self.table}" SET "models" = array_append("models", $2) ' - f'WHERE $1 = ANY("models") AND NOT ($2 = ANY("models")) RETURNING {self.returning}' + f'{self.kind}_rows AS (UPDATE "{self.table}" SET "models" = {set_clause} WHERE {where_clause} ' + f"RETURNING '{self.kind}' AS kind, \"{self.id_column}\" AS object_id, {alias} AS team_alias)" ) @@ -63,27 +62,28 @@ def _user_cache_keys(row: _TouchedRow) -> tuple[str, ...]: _ALLOWLIST_TABLES: Final = ( - _AllowlistTable("LiteLLM_TeamTable", '"team_id" AS object_id, "team_alias"', _team_cache_keys), - _AllowlistTable("LiteLLM_VerificationToken", '"token" AS object_id', _key_cache_keys), - _AllowlistTable("LiteLLM_OrganizationTable", '"organization_id" AS object_id', _org_cache_keys), - _AllowlistTable("LiteLLM_ProjectTable", '"project_id" AS object_id', _project_cache_keys), - _AllowlistTable("LiteLLM_UserTable", '"user_id" AS object_id', _user_cache_keys), + _AllowlistTable("team", "LiteLLM_TeamTable", "team_id", _team_cache_keys, alias_column="team_alias"), + _AllowlistTable("key", "LiteLLM_VerificationToken", "token", _key_cache_keys), + _AllowlistTable("org", "LiteLLM_OrganizationTable", "organization_id", _org_cache_keys), + _AllowlistTable("project", "LiteLLM_ProjectTable", "project_id", _project_cache_keys), + _AllowlistTable("user", "LiteLLM_UserTable", "user_id", _user_cache_keys), ) +_CACHE_KEYS_BY_KIND: Final = MappingProxyType({table.kind: table.cache_keys for table in _ALLOWLIST_TABLES}) -async def _rewrite_allowlist( - executor: RawExecutor, - allowlist: _AllowlistTable, - sql: str, - old_name: str, - new_name: str, - user_api_key_cache: UserApiKeyCache, -) -> None: - touched_rows: Final = await executor.query_raw(sql, old_name, new_name) - await evict_and_broadcast( - tuple(cache_key for row in touched_rows for cache_key in allowlist.cache_keys(_TouchedRow.model_validate(row))), - user_api_key_cache, + +def _rewrite_sql(set_clause: str, where_clause: str) -> str: + """One statement touching every allowlist table, so the rewrite lands everywhere or nowhere.""" + ctes: Final = ", ".join(table.update_cte(set_clause, where_clause) for table in _ALLOWLIST_TABLES) + rows: Final = " UNION ALL ".join( + f"SELECT kind, object_id, team_alias FROM {table.kind}_rows" for table in _ALLOWLIST_TABLES ) + return f"WITH {ctes} {rows}" + + +_REPLACE_SQL: Final = _rewrite_sql('array_replace(array_remove("models", $2), $1, $2)', '$1 = ANY("models")') + +_APPEND_SQL: Final = _rewrite_sql('array_append("models", $2)', '$1 = ANY("models") AND NOT ($2 = ANY("models"))') async def sync_model_allowlists_for_renamed_model( @@ -99,12 +99,11 @@ async def sync_model_allowlists_for_renamed_model( return executor: Final = raw_executor(prisma_client) old_name_still_backed: Final = await still_backed(executor, llm_router, old_name, model_id) - for allowlist in _ALLOWLIST_TABLES: - await _rewrite_allowlist( - executor, - allowlist, - allowlist.append_sql() if old_name_still_backed else allowlist.replace_sql(), - old_name, - new_name, - user_api_key_cache, - ) + touched_rows: Final = await executor.query_raw( + _APPEND_SQL if old_name_still_backed else _REPLACE_SQL, old_name, new_name + ) + touched: Final = tuple(_TouchedRow.model_validate(row) for row in touched_rows) + await evict_and_broadcast( + tuple(cache_key for row in touched for cache_key in _CACHE_KEYS_BY_KIND[row.kind](row)), + user_api_key_cache, + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 0f273ebce84..20e51e7c906 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -6064,13 +6064,21 @@ class TestAccessGroupModelSync: _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" _INVALIDATE = "litellm.proxy.management_helpers.access_group_model_sync.invalidate_access_group_caches" _EVICT = "litellm.proxy.management_helpers.model_allowlist_rename_sync.evict_and_broadcast" - _ALLOWLIST_ROWS = { - "LiteLLM_TeamTable": [{"object_id": "team-1", "team_alias": "alias-1"}, {"object_id": "team-2", "team_alias": None}], - "LiteLLM_VerificationToken": [{"object_id": "hashed-token-1"}], - "LiteLLM_OrganizationTable": [{"object_id": "org-1"}], - "LiteLLM_ProjectTable": [{"object_id": "proj-1"}], - "LiteLLM_UserTable": [{"object_id": "user-1"}], - } + _ALLOWLIST_TABLES = ( + "LiteLLM_TeamTable", + "LiteLLM_VerificationToken", + "LiteLLM_OrganizationTable", + "LiteLLM_ProjectTable", + "LiteLLM_UserTable", + ) + _ALLOWLIST_ROWS = [ + {"kind": "team", "object_id": "team-1", "team_alias": "alias-1"}, + {"kind": "team", "object_id": "team-2", "team_alias": None}, + {"kind": "key", "object_id": "hashed-token-1", "team_alias": None}, + {"kind": "org", "object_id": "org-1", "team_alias": None}, + {"kind": "project", "object_id": "proj-1", "team_alias": None}, + {"kind": "user", "object_id": "user-1", "team_alias": None}, + ] @staticmethod def _admin(): @@ -6092,7 +6100,8 @@ class TestAccessGroupModelSync: return [{"deployment_count": deployment_count}] if sql.startswith('UPDATE "LiteLLM_AccessGroupTable"'): return [{"access_group_id": "ag-1"}] - return TestAccessGroupModelSync._ALLOWLIST_ROWS[sql.split('"')[1]] + assert sql.startswith("WITH ") + return TestAccessGroupModelSync._ALLOWLIST_ROWS mock_prisma = MagicMock() mock_prisma.db = MagicMock() @@ -6113,11 +6122,11 @@ class TestAccessGroupModelSync: @staticmethod def _allowlist_updates(mock_prisma): - return { - call.args[0].split('"')[1]: call + return [ + call for call in mock_prisma.db.query_raw.await_args_list - if call.args[0].startswith('UPDATE "') and 'SET "models"' in call.args[0] - } + if call.args[0].startswith("WITH ") and 'SET "models"' in call.args[0] + ] @contextlib.contextmanager def _endpoint_env(self, mock_prisma, router, evict=None): @@ -6130,7 +6139,9 @@ class TestAccessGroupModelSync: patch(f"{self._PS}.proxy_logging_obj", MagicMock()), patch(f"{self._PS}.user_api_key_cache", MagicMock()), patch(self._EVICT, new=evict or AsyncMock()), - patch(f"{self._MOD}.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None)), + patch( + f"{self._MOD}.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None) + ), patch( f"{self._MOD}.clear_cache", new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), @@ -6190,7 +6201,9 @@ class TestAccessGroupModelSync: router.get_model_ids.return_value = ["m-same"] with self._endpoint_env(mock_prisma, router) as invalidate: - await patch_model(model_id="m-same", patch_data=updateDeployment(blocked=True), user_api_key_dict=self._admin()) + await patch_model( + model_id="m-same", patch_data=updateDeployment(blocked=True), user_api_key_dict=self._admin() + ) mock_prisma.db.query_raw.assert_not_awaited() invalidate.assert_not_awaited() @@ -6278,20 +6291,24 @@ class TestAccessGroupModelSync: user_api_key_dict=self._admin(), ) - updates = self._allowlist_updates(mock_prisma) - assert set(updates) == set(self._ALLOWLIST_ROWS) - for update_call in updates.values(): - assert 'SET "models" = array_replace(array_remove("models", $2), $1, $2)' in update_call.args[0] - assert 'WHERE $1 = ANY("models")' in update_call.args[0] - assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") - evicted = [call.args[0] for call in evict.await_args_list] - assert evicted == [ - ("team_id:team-1", "team_alias:alias-1", "team_id:team-2"), - ("hashed-token-1",), - ("org_id:org-1", "org_id:org-1:with_budget"), - ("project_id:proj-1",), - ("user-1",), - ] + (update_call,) = self._allowlist_updates(mock_prisma) + for table in self._ALLOWLIST_TABLES: + assert ( + f'UPDATE "{table}" SET "models" = array_replace(array_remove("models", $2), $1, $2) ' + 'WHERE $1 = ANY("models") RETURNING' + ) in update_call.args[0] + assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") + evict.assert_awaited_once() + assert evict.await_args.args[0] == ( + "team_id:team-1", + "team_alias:alias-1", + "team_id:team-2", + "hashed-token-1", + "org_id:org-1", + "org_id:org-1:with_budget", + "project_id:proj-1", + "user-1", + ) @pytest.mark.asyncio async def test_rename_appends_to_allowlists_when_a_sibling_deployment_keeps_the_old_name(self): @@ -6308,12 +6325,13 @@ class TestAccessGroupModelSync: user_api_key_dict=self._admin(), ) - updates = self._allowlist_updates(mock_prisma) - assert set(updates) == set(self._ALLOWLIST_ROWS) - for update_call in updates.values(): - assert 'SET "models" = array_append("models", $2)' in update_call.args[0] - assert 'WHERE $1 = ANY("models") AND NOT ($2 = ANY("models"))' in update_call.args[0] - assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") + (update_call,) = self._allowlist_updates(mock_prisma) + for table in self._ALLOWLIST_TABLES: + assert ( + f'UPDATE "{table}" SET "models" = array_append("models", $2) ' + 'WHERE $1 = ANY("models") AND NOT ($2 = ANY("models")) RETURNING' + ) in update_call.args[0] + assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") @pytest.mark.asyncio async def test_unchanged_name_never_touches_allowlists(self): @@ -6334,7 +6352,7 @@ class TestAccessGroupModelSync: user_api_key_cache=MagicMock(), ) - assert self._allowlist_updates(mock_prisma) == {} + assert self._allowlist_updates(mock_prisma) == [] evict.assert_not_awaited()