From 050b4f78bbd65ddfcb3cab71ec2f92998bbd4e76 Mon Sep 17 00:00:00 2001 From: RoyVivat Date: Wed, 8 Apr 2026 16:00:16 -0700 Subject: [PATCH] fix(access_group): eliminate N+1 queries, fix UUID persistence, and break cyclic imports - Batch-fetch teams/keys/object-permissions before loops in all four sync helpers - Include object_permission_id in upsert create payload so the generated UUID is stored - Move _merge_access_group_resources_into_data_json import to lazy (inside callers) in team/key endpoints to break the module-level cyclic import flagged by CodeQL - Remove unused Tuple import Co-Authored-By: Claude Sonnet 4.6 --- .../access_group_endpoints.py | 185 +++++++++++++++--- .../key_management_endpoints.py | 10 +- .../management_endpoints/team_endpoints.py | 11 +- 3 files changed, 172 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 17d281d8d1c..f1889821a8e 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Set, Tuple +from typing import Dict, List, Optional, Set from fastapi import APIRouter, Depends, HTTPException, status @@ -156,6 +156,7 @@ async def _upsert_mcp_agents_in_object_permission( existing_op_id: Optional[str], ag_mcp_servers: List[str], ag_agents: List[str], + existing_op=None, ) -> Optional[str]: """ Upsert LiteLLM_ObjectPermissionTable to add MCP servers and agents. @@ -163,6 +164,8 @@ async def _upsert_mcp_agents_in_object_permission( Merges ``ag_mcp_servers`` / ``ag_agents`` into the existing record (if any), creating a new record when ``existing_op_id`` is None. + ``existing_op``: optional pre-fetched record to avoid an extra DB round-trip. + Returns the ``object_permission_id`` of the upserted row, or ``None`` when both lists are empty (nothing to do). """ @@ -174,9 +177,10 @@ async def _upsert_mcp_agents_in_object_permission( existing_data: Dict = {} if existing_op_id: - existing_op = await tx.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": existing_op_id} - ) + if existing_op is None: + existing_op = await tx.litellm_objectpermissiontable.find_unique( + where={"object_permission_id": existing_op_id} + ) if existing_op is not None: try: existing_data = existing_op.model_dump(exclude_none=True) @@ -194,9 +198,10 @@ async def _upsert_mcp_agents_in_object_permission( upsert_data["agents"] = list(set(existing_agents + ag_agents)) op_id_to_use: str = existing_op_id or str(uuid.uuid4()) + create_data: Dict = {**upsert_data, "object_permission_id": op_id_to_use} created_row = await tx.litellm_objectpermissiontable.upsert( where={"object_permission_id": op_id_to_use}, - data={"create": upsert_data, "update": upsert_data}, + data={"create": create_data, "update": upsert_data}, ) return created_row.object_permission_id @@ -206,18 +211,22 @@ async def _remove_mcp_agents_from_object_permission( existing_op_id: Optional[str], mcp_servers_to_remove: List[str], agents_to_remove: List[str], + existing_op=None, ) -> None: """ Remove specific MCP server IDs and agent IDs from an existing LiteLLM_ObjectPermissionTable row. No-ops when the record does not exist or the removal sets are empty. + + ``existing_op``: optional pre-fetched record to avoid an extra DB round-trip. """ if not existing_op_id or (not mcp_servers_to_remove and not agents_to_remove): return - existing_op = await tx.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": existing_op_id} - ) + if existing_op is None: + existing_op = await tx.litellm_objectpermissiontable.find_unique( + where={"object_permission_id": existing_op_id} + ) if existing_op is None: return @@ -268,8 +277,32 @@ async def _sync_add_access_group_to_teams( getattr(access_group_record, "access_agent_ids", None) or [] ) + if not team_ids: + return + + # Batch-fetch all teams to avoid N+1 queries. + teams = await tx.litellm_teamtable.find_many( + where={"team_id": {"in": team_ids}} + ) + team_map: Dict = {t.team_id: t for t in teams} + + # Batch-fetch object permissions for teams that need MCP/agent merging. + op_map: Dict = {} + if ag_mcp_servers or ag_agents: + op_ids = [ + t.object_permission_id + for t in teams + if getattr(t, "object_permission_id", None) + and access_group_id not in (t.access_group_ids or []) + ] + if op_ids: + op_records = await tx.litellm_objectpermissiontable.find_many( + where={"object_permission_id": {"in": op_ids}} + ) + op_map = {r.object_permission_id: r for r in op_records} + for team_id in team_ids: - team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) + team = team_map.get(team_id) if team is None or access_group_id in (team.access_group_ids or []): continue @@ -292,6 +325,7 @@ async def _sync_add_access_group_to_teams( existing_op_id=existing_op_id, ag_mcp_servers=ag_mcp_servers, ag_agents=ag_agents, + existing_op=op_map.get(existing_op_id) if existing_op_id else None, ) # Link the (possibly newly created) object_permission row to the team if new_op_id is not None and new_op_id != existing_op_id: @@ -336,19 +370,53 @@ async def _sync_remove_access_group_from_teams( getattr(ag_record, "access_agent_ids", None) or [] ) - for team_id in team_ids: - team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) - if team is None or access_group_id not in (team.access_group_ids or []): - continue + if not team_ids: + return + # Batch-fetch all teams to avoid N+1 queries. + teams = await tx.litellm_teamtable.find_many( + where={"team_id": {"in": team_ids}} + ) + relevant_teams = [ + t for t in teams if access_group_id in (t.access_group_ids or []) + ] + + # Collect all unique remaining AG IDs across affected teams so we can + # batch-fetch their records in one query instead of one per team. + all_remaining_ag_ids: Set[str] = set() + for team in relevant_teams: + for ag in (team.access_group_ids or []): + if ag != access_group_id: + all_remaining_ag_ids.add(ag) + + remaining_ag_map: Dict = {} + if all_remaining_ag_ids: + remaining_ag_records = await tx.litellm_accessgrouptable.find_many( + where={"access_group_id": {"in": list(all_remaining_ag_ids)}} + ) + remaining_ag_map = {r.access_group_id: r for r in remaining_ag_records} + + # Batch-fetch object permissions for affected teams. + all_op_ids = [ + t.object_permission_id + for t in relevant_teams + if getattr(t, "object_permission_id", None) + ] + op_map: Dict = {} + if all_op_ids: + op_records = await tx.litellm_objectpermissiontable.find_many( + where={"object_permission_id": {"in": all_op_ids}} + ) + op_map = {r.object_permission_id: r for r in op_records} + + for team in relevant_teams: remaining_ag_ids = [ ag for ag in (team.access_group_ids or []) if ag != access_group_id ] + remaining_records = [ + remaining_ag_map[ag] for ag in remaining_ag_ids if ag in remaining_ag_map + ] - # Batch-fetch remaining AGs to compute what resources they still provide - remaining_records = await tx.litellm_accessgrouptable.find_many( - where={"access_group_id": {"in": remaining_ag_ids}} - ) models_in_remaining: Set[str] = { m for r in remaining_records for m in (r.access_model_names or []) } @@ -378,10 +446,11 @@ async def _sync_remove_access_group_from_teams( existing_op_id=existing_op_id, mcp_servers_to_remove=mcp_to_remove, agents_to_remove=agents_to_remove, + existing_op=op_map.get(existing_op_id) if existing_op_id else None, ) await tx.litellm_teamtable.update( - where={"team_id": team_id}, + where={"team_id": team.team_id}, data=update_data, ) @@ -402,8 +471,32 @@ async def _sync_add_access_group_to_keys( getattr(access_group_record, "access_agent_ids", None) or [] ) + if not key_tokens: + return + + # Batch-fetch all keys to avoid N+1 queries. + keys = await tx.litellm_verificationtoken.find_many( + where={"token": {"in": key_tokens}} + ) + key_map: Dict = {k.token: k for k in keys} + + # Batch-fetch object permissions for keys that need MCP/agent merging. + op_map: Dict = {} + if ag_mcp_servers or ag_agents: + op_ids = [ + k.object_permission_id + for k in keys + if getattr(k, "object_permission_id", None) + and access_group_id not in (k.access_group_ids or []) + ] + if op_ids: + op_records = await tx.litellm_objectpermissiontable.find_many( + where={"object_permission_id": {"in": op_ids}} + ) + op_map = {r.object_permission_id: r for r in op_records} + for token in key_tokens: - key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) + key = key_map.get(token) if key is None or access_group_id in (key.access_group_ids or []): continue @@ -425,6 +518,7 @@ async def _sync_add_access_group_to_keys( existing_op_id=existing_op_id, ag_mcp_servers=ag_mcp_servers, ag_agents=ag_agents, + existing_op=op_map.get(existing_op_id) if existing_op_id else None, ) # Link the (possibly newly created) object_permission row to the key if new_op_id is not None and new_op_id != existing_op_id: @@ -469,19 +563,53 @@ async def _sync_remove_access_group_from_keys( getattr(ag_record, "access_agent_ids", None) or [] ) - for token in key_tokens: - key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) - if key is None or access_group_id not in (key.access_group_ids or []): - continue + if not key_tokens: + return + # Batch-fetch all keys to avoid N+1 queries. + keys = await tx.litellm_verificationtoken.find_many( + where={"token": {"in": key_tokens}} + ) + relevant_keys = [ + k for k in keys if access_group_id in (k.access_group_ids or []) + ] + + # Collect all unique remaining AG IDs across affected keys so we can + # batch-fetch their records in one query instead of one per key. + all_remaining_ag_ids: Set[str] = set() + for key in relevant_keys: + for ag in (key.access_group_ids or []): + if ag != access_group_id: + all_remaining_ag_ids.add(ag) + + remaining_ag_map: Dict = {} + if all_remaining_ag_ids: + remaining_ag_records = await tx.litellm_accessgrouptable.find_many( + where={"access_group_id": {"in": list(all_remaining_ag_ids)}} + ) + remaining_ag_map = {r.access_group_id: r for r in remaining_ag_records} + + # Batch-fetch object permissions for affected keys. + all_op_ids = [ + k.object_permission_id + for k in relevant_keys + if getattr(k, "object_permission_id", None) + ] + op_map: Dict = {} + if all_op_ids: + op_records = await tx.litellm_objectpermissiontable.find_many( + where={"object_permission_id": {"in": all_op_ids}} + ) + op_map = {r.object_permission_id: r for r in op_records} + + for key in relevant_keys: remaining_ag_ids = [ ag for ag in (key.access_group_ids or []) if ag != access_group_id ] + remaining_records = [ + remaining_ag_map[ag] for ag in remaining_ag_ids if ag in remaining_ag_map + ] - # Batch-fetch remaining AGs to compute what resources they still provide - remaining_records = await tx.litellm_accessgrouptable.find_many( - where={"access_group_id": {"in": remaining_ag_ids}} - ) models_in_remaining: Set[str] = { m for r in remaining_records for m in (r.access_model_names or []) } @@ -506,10 +634,11 @@ async def _sync_remove_access_group_from_keys( existing_op_id=existing_op_id, mcp_servers_to_remove=mcp_to_remove, agents_to_remove=agents_to_remove, + existing_op=op_map.get(existing_op_id) if existing_op_id else None, ) await tx.litellm_verificationtoken.update( - where={"token": token}, + where={"token": key.token}, data={ "access_group_ids": remaining_ag_ids, "models": updated_models, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2fa6a55945d..c588d5a9bec 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -56,9 +56,6 @@ from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _set_object_metadata_field, ) -from litellm.proxy.management_endpoints.access_group_endpoints import ( - _merge_access_group_resources_into_data_json, -) from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) @@ -678,6 +675,10 @@ async def _common_key_generation_helper( # noqa: PLR0915 # Populate models/MCP servers/agents from key-level access groups (if provided). # Only runs when access_group_ids is explicitly included in the request. if data_json.get("access_group_ids"): + from litellm.proxy.management_endpoints.access_group_endpoints import ( + _merge_access_group_resources_into_data_json, + ) + data_json = await _merge_access_group_resources_into_data_json( data_json=data_json, access_group_ids=data_json["access_group_ids"], @@ -1576,6 +1577,9 @@ async def prepare_key_update_data( if "access_group_ids" in non_default_values: new_access_group_ids = non_default_values.get("access_group_ids") or [] if new_access_group_ids: + from litellm.proxy.management_endpoints.access_group_endpoints import ( + _merge_access_group_resources_into_data_json, + ) from litellm.proxy.proxy_server import prisma_client as _prisma_client if _prisma_client is not None: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 619c0b539ca..8d9118ac2b7 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -78,9 +78,6 @@ from litellm.proxy.management_endpoints.common_utils import ( _upsert_budget_and_membership, _user_has_admin_view, ) -from litellm.proxy.management_endpoints.access_group_endpoints import ( - _merge_access_group_resources_into_data_json, -) from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, ) @@ -933,6 +930,10 @@ async def new_team( # noqa: PLR0915 ## so they appear in model lists and MCP validation at key-generation time. ## Only runs when access_group_ids is explicitly provided in this request. if data.access_group_ids: + from litellm.proxy.management_endpoints.access_group_endpoints import ( + _merge_access_group_resources_into_data_json, + ) + data_json = await _merge_access_group_resources_into_data_json( data_json=data_json, access_group_ids=data.access_group_ids, @@ -1534,6 +1535,10 @@ async def update_team( # noqa: PLR0915 if "access_group_ids" in updated_kv: new_access_group_ids = updated_kv.get("access_group_ids") or [] if new_access_group_ids: + from litellm.proxy.management_endpoints.access_group_endpoints import ( + _merge_access_group_resources_into_data_json, + ) + updated_kv = await _merge_access_group_resources_into_data_json( data_json=updated_kv, access_group_ids=new_access_group_ids,