mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
refactor(proxy): extract helpers to fix PLR0915 violations
Extract `_apply_non_admin_alias_scope` from `key_aliases`, `_resolve_team_access_group_resources` from `team_info`, and `_enforce_list_team_v2_access` from `list_team_v2` to bring each function under ruff's 50-statement limit. No behavior changes.
This commit is contained in:
parent
0331fb5a8f
commit
ce219fcc96
2 changed files with 134 additions and 85 deletions
|
|
@ -4382,6 +4382,42 @@ async def list_keys(
|
|||
)
|
||||
|
||||
|
||||
async def _apply_non_admin_alias_scope(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: Any,
|
||||
query_params: List[Any],
|
||||
where_parts: List[str],
|
||||
) -> None:
|
||||
"""Append SQL scope conditions so non-admin users only see aliases for
|
||||
keys they own or keys belonging to teams they are members of."""
|
||||
scope_conditions: List[str] = []
|
||||
if user_api_key_dict.user_id:
|
||||
query_params.append(user_api_key_dict.user_id)
|
||||
scope_conditions.append(f"user_id = ${len(query_params)}")
|
||||
|
||||
# Look up the user's teams from the user table
|
||||
user_teams: List[str] = []
|
||||
if user_api_key_dict.user_id:
|
||||
user_row = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_api_key_dict.user_id}
|
||||
)
|
||||
if user_row is not None:
|
||||
user_teams = getattr(user_row, "teams", []) or []
|
||||
|
||||
if user_teams:
|
||||
team_placeholders = ", ".join(
|
||||
f"${len(query_params) + i + 1}" for i in range(len(user_teams))
|
||||
)
|
||||
query_params.extend(user_teams)
|
||||
scope_conditions.append(f"team_id IN ({team_placeholders})")
|
||||
|
||||
if scope_conditions:
|
||||
where_parts.append(f"({' OR '.join(scope_conditions)})")
|
||||
else:
|
||||
# No user_id and no teams — return nothing
|
||||
where_parts.append("FALSE")
|
||||
|
||||
|
||||
@router.get(
|
||||
"/key/aliases",
|
||||
tags=["key management"],
|
||||
|
|
@ -4442,32 +4478,9 @@ async def key_aliases(
|
|||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
|
||||
]
|
||||
if not is_proxy_admin:
|
||||
scope_conditions: List[str] = []
|
||||
if user_api_key_dict.user_id:
|
||||
query_params.append(user_api_key_dict.user_id)
|
||||
scope_conditions.append(f"user_id = ${len(query_params)}")
|
||||
|
||||
# Look up the user's teams from the user table
|
||||
user_teams: List[str] = []
|
||||
if user_api_key_dict.user_id:
|
||||
user_row = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_api_key_dict.user_id}
|
||||
)
|
||||
if user_row is not None:
|
||||
user_teams = getattr(user_row, "teams", []) or []
|
||||
|
||||
if user_teams:
|
||||
team_placeholders = ", ".join(
|
||||
f"${len(query_params) + i + 1}" for i in range(len(user_teams))
|
||||
)
|
||||
query_params.extend(user_teams)
|
||||
scope_conditions.append(f"team_id IN ({team_placeholders})")
|
||||
|
||||
if scope_conditions:
|
||||
where_parts.append(f"({' OR '.join(scope_conditions)})")
|
||||
else:
|
||||
# No user_id and no teams — return nothing
|
||||
where_parts.append("FALSE")
|
||||
await _apply_non_admin_alias_scope(
|
||||
user_api_key_dict, prisma_client, query_params, where_parts
|
||||
)
|
||||
|
||||
if search:
|
||||
query_params.append(f"%{search}%")
|
||||
|
|
|
|||
|
|
@ -2932,6 +2932,25 @@ async def _add_team_member_budget_table(
|
|||
return team_info_response_object
|
||||
|
||||
|
||||
async def _resolve_team_access_group_resources(_team_info: Any) -> None:
|
||||
"""Populate access_group_models / mcp_server_ids / agent_ids on the team
|
||||
info response by resolving inherited resources from its access groups."""
|
||||
if not _team_info.access_group_ids:
|
||||
return
|
||||
ag_lookup = await _batch_resolve_access_group_resources(
|
||||
_team_info.access_group_ids
|
||||
)
|
||||
models, mcp_ids, agent_ids = set(), set(), set()
|
||||
for ag_id in _team_info.access_group_ids:
|
||||
if ag_id in ag_lookup:
|
||||
models.update(ag_lookup[ag_id]["models"])
|
||||
mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"])
|
||||
agent_ids.update(ag_lookup[ag_id]["agent_ids"])
|
||||
_team_info.access_group_models = list(models)
|
||||
_team_info.access_group_mcp_server_ids = list(mcp_ids)
|
||||
_team_info.access_group_agent_ids = list(agent_ids)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
|
||||
)
|
||||
|
|
@ -3043,17 +3062,7 @@ async def team_info(
|
|||
)
|
||||
|
||||
# Resolve resources inherited from access groups
|
||||
if _team_info.access_group_ids:
|
||||
ag_lookup = await _batch_resolve_access_group_resources(_team_info.access_group_ids)
|
||||
models, mcp_ids, agent_ids = set(), set(), set()
|
||||
for ag_id in _team_info.access_group_ids:
|
||||
if ag_id in ag_lookup:
|
||||
models.update(ag_lookup[ag_id]["models"])
|
||||
mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"])
|
||||
agent_ids.update(ag_lookup[ag_id]["agent_ids"])
|
||||
_team_info.access_group_models = list(models)
|
||||
_team_info.access_group_mcp_server_ids = list(mcp_ids)
|
||||
_team_info.access_group_agent_ids = list(agent_ids)
|
||||
await _resolve_team_access_group_resources(_team_info)
|
||||
|
||||
response_object = TeamInfoResponseObject(
|
||||
team_id=team_id,
|
||||
|
|
@ -3401,6 +3410,73 @@ def _convert_teams_to_response_models(
|
|||
return team_list
|
||||
|
||||
|
||||
async def _enforce_list_team_v2_access(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
user_id: Optional[str],
|
||||
organization_id: Optional[str],
|
||||
prisma_client: Any,
|
||||
user_api_key_cache: Any,
|
||||
proxy_logging_obj: Any,
|
||||
) -> Tuple[Optional[str], Optional[List[str]]]:
|
||||
"""Enforce access control for list_team_v2.
|
||||
|
||||
- Proxy admins and admin viewers can query any teams.
|
||||
- Org admins can query teams within their organizations.
|
||||
- Regular users can only query their own teams.
|
||||
|
||||
Returns the (possibly overridden) user_id and org_admin_org_ids.
|
||||
"""
|
||||
is_proxy_admin = _user_has_admin_view(user_api_key_dict)
|
||||
org_admin_org_ids: Optional[List[str]] = None
|
||||
|
||||
if is_proxy_admin:
|
||||
return user_id, org_admin_org_ids
|
||||
|
||||
# Always check org admin status so that even own-queries see
|
||||
# the full set of organisation teams, not just direct memberships.
|
||||
if user_api_key_dict.user_id:
|
||||
org_admin_org_ids = await _get_org_admin_org_ids(
|
||||
user_id=user_api_key_dict.user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
if org_admin_org_ids is not None:
|
||||
# Org admin: validate org_id filter if provided
|
||||
if organization_id and organization_id not in org_admin_org_ids:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "You can only view teams within your organizations."
|
||||
},
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"list_team_v2: org admin access for user=%s, org_ids=%s, user_id_filter=%s",
|
||||
user_api_key_dict.user_id,
|
||||
org_admin_org_ids,
|
||||
user_id,
|
||||
)
|
||||
else:
|
||||
# Not an org admin — fall back to standard route check
|
||||
if not allowed_route_check_inside_route(
|
||||
user_api_key_dict=user_api_key_dict, requested_user_id=user_id
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": "Only admin users can query all teams/other teams. Your user role={}".format(
|
||||
user_api_key_dict.user_role
|
||||
)
|
||||
},
|
||||
)
|
||||
# Regular user — auto-inject caller's user_id
|
||||
if user_id is None:
|
||||
user_id = user_api_key_dict.user_id
|
||||
|
||||
return user_id, org_admin_org_ids
|
||||
|
||||
|
||||
@router.get(
|
||||
"/v2/team/list",
|
||||
tags=["team management"],
|
||||
|
|
@ -3478,54 +3554,14 @@ async def list_team_v2(
|
|||
)
|
||||
|
||||
# --- Access control ---
|
||||
# Proxy admins and admin viewers can query any teams.
|
||||
# Org admins can query teams within their organizations.
|
||||
# Regular users can only query their own teams.
|
||||
is_proxy_admin = _user_has_admin_view(user_api_key_dict)
|
||||
org_admin_org_ids: Optional[List[str]] = None
|
||||
|
||||
if not is_proxy_admin:
|
||||
# Always check org admin status so that even own-queries see
|
||||
# the full set of organisation teams, not just direct memberships.
|
||||
if user_api_key_dict.user_id:
|
||||
org_admin_org_ids = await _get_org_admin_org_ids(
|
||||
user_id=user_api_key_dict.user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
if org_admin_org_ids is not None:
|
||||
# Org admin: validate org_id filter if provided
|
||||
if organization_id and organization_id not in org_admin_org_ids:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "You can only view teams within your organizations."
|
||||
},
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"list_team_v2: org admin access for user=%s, org_ids=%s, user_id_filter=%s",
|
||||
user_api_key_dict.user_id,
|
||||
org_admin_org_ids,
|
||||
user_id,
|
||||
)
|
||||
else:
|
||||
# Not an org admin — fall back to standard route check
|
||||
if not allowed_route_check_inside_route(
|
||||
user_api_key_dict=user_api_key_dict, requested_user_id=user_id
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": "Only admin users can query all teams/other teams. Your user role={}".format(
|
||||
user_api_key_dict.user_role
|
||||
)
|
||||
},
|
||||
)
|
||||
# Regular user — auto-inject caller's user_id
|
||||
if user_id is None:
|
||||
user_id = user_api_key_dict.user_id
|
||||
user_id, org_admin_org_ids = await _enforce_list_team_v2_access(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
user_id=user_id,
|
||||
organization_id=organization_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
if status is not None and status != "deleted":
|
||||
raise HTTPException(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue