From 1d403e9bd8141fb34f6b37bea95a7e6d4d6b30e3 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 16:29:28 -0700 Subject: [PATCH] refactor(key_management): extract validation logic from update_key_fn to fix PLR0915 Extract permission checks and constraint validation from update_key_fn into _validate_update_key_data helper to reduce statement count below the 50-statement limit. Co-Authored-By: Claude Opus 4.6 --- .../key_management_endpoints.py | 255 ++++++++++-------- 1 file changed, 136 insertions(+), 119 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4101440964b..db1a089ff7b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1799,6 +1799,139 @@ async def _validate_mcp_servers_for_key_update( ) +async def _validate_update_key_data( + data: UpdateKeyRequest, + existing_key_row: Any, + user_api_key_dict: UserAPIKeyAuth, + llm_router: Any, + premium_user: bool, + prisma_client: Any, + user_api_key_cache: Any, +) -> None: + """Validate permissions and constraints for key update.""" + # sanity check - prevent non-proxy admin user from updating key to belong to a different user + if ( + data.user_id is not None + and data.user_id != existing_key_row.user_id + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value + ): + raise HTTPException( + status_code=403, + detail=f"User={data.user_id} is not allowed to update key={data.key} to belong to user={existing_key_row.user_id}", + ) + + common_key_access_checks( + user_api_key_dict=user_api_key_dict, + data=data, + user_id=existing_key_row.user_id, + llm_router=llm_router, + premium_user=premium_user, + ) + + await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_UPDATE, + prisma_client=prisma_client, + existing_key_row=existing_key_row, + user_api_key_cache=user_api_key_cache, + ) + + # Check team limits if key has a team_id (from request or existing key) + team_obj: Optional[LiteLLM_TeamTableCachedObj] = None + _team_id_to_check = data.team_id or getattr( + existing_key_row, "team_id", None + ) + if _team_id_to_check is not None: + team_obj = await get_team_object( + team_id=_team_id_to_check, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + + if team_obj is not None: + await _check_team_key_limits( + team_table=team_obj, + data=data, + prisma_client=prisma_client, + ) + + # Validate key against project limits if project_id is being set + _project_id_to_check = getattr(data, "project_id", None) or getattr( + existing_key_row, "project_id", None + ) + if _project_id_to_check is not None and ( + data.models is not None or data.max_budget is not None + ): + await _check_project_key_limits( + project_id=_project_id_to_check, + data=data, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + # Check org key limits only when throughput-related fields or organization_id change + _org_id_to_check = data.organization_id or getattr( + existing_key_row, "organization_id", None + ) + _throughput_fields_changed = ( + data.organization_id is not None + or data.tpm_limit is not None + or data.rpm_limit is not None + or data.tpm_limit_type is not None + or data.rpm_limit_type is not None + ) + if _org_id_to_check is not None and _throughput_fields_changed: + org_table = await get_org_object( + org_id=_org_id_to_check, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + ) + if org_table is None: + raise HTTPException( + status_code=400, + detail=f"Organization not found for organization_id={_org_id_to_check}", + ) + await _check_org_key_limits( + org_table=org_table, + data=data, + prisma_client=prisma_client, + ) + + # if team change - check if this is possible + if is_different_team(data=data, existing_key_row=existing_key_row): + if llm_router is None: + raise HTTPException( + status_code=400, + detail={ + "error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI." + }, + ) + if team_obj is None: + raise HTTPException( + status_code=500, + detail={ + "error": "Team object not found for team change validation" + }, + ) + await validate_key_team_change( + key=existing_key_row, + team=team_obj, + change_initiated_by=user_api_key_dict, + llm_router=llm_router, + ) + + # Validate MCP servers in object_permission against the effective team + if data.object_permission is not None: + await _validate_mcp_servers_for_key_update( + data=data, + team_obj=team_obj, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + @router.post( "/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) @@ -1913,132 +2046,16 @@ async def update_key_fn( detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) - ## sanity check - prevent non-proxy admin user from updating key to belong to a different user - if ( - data.user_id is not None - and data.user_id != existing_key_row.user_id - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - ): - raise HTTPException( - status_code=403, - detail=f"User={data.user_id} is not allowed to update key={key} to belong to user={existing_key_row.user_id}", - ) - - common_key_access_checks( - user_api_key_dict=user_api_key_dict, + await _validate_update_key_data( data=data, - user_id=existing_key_row.user_id, + existing_key_row=existing_key_row, + user_api_key_dict=user_api_key_dict, llm_router=llm_router, premium_user=premium_user, - ) - - # check if user has permission to update key - await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_UPDATE, prisma_client=prisma_client, - existing_key_row=existing_key_row, user_api_key_cache=user_api_key_cache, ) - # Check team limits if key has a team_id (from request or existing key) - team_obj: Optional[LiteLLM_TeamTableCachedObj] = None - _team_id_to_check = data.team_id or getattr( - existing_key_row, "team_id", None - ) - if _team_id_to_check is not None: - team_obj = await get_team_object( - team_id=_team_id_to_check, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - - if team_obj is not None: - await _check_team_key_limits( - team_table=team_obj, - data=data, - prisma_client=prisma_client, - ) - - # Validate key against project limits if project_id is being set - _project_id_to_check = getattr(data, "project_id", None) or getattr( - existing_key_row, "project_id", None - ) - if _project_id_to_check is not None and ( - data.models is not None or data.max_budget is not None - ): - await _check_project_key_limits( - project_id=_project_id_to_check, - data=data, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - - # Check org key limits only when throughput-related fields or organization_id change - _org_id_to_check = data.organization_id or getattr( - existing_key_row, "organization_id", None - ) - _throughput_fields_changed = ( - data.organization_id is not None - or data.tpm_limit is not None - or data.rpm_limit is not None - or data.tpm_limit_type is not None - or data.rpm_limit_type is not None - ) - if _org_id_to_check is not None and _throughput_fields_changed: - org_table = await get_org_object( - org_id=_org_id_to_check, - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, - ) - if org_table is None: - raise HTTPException( - status_code=400, - detail=f"Organization not found for organization_id={_org_id_to_check}", - ) - await _check_org_key_limits( - org_table=org_table, - data=data, - prisma_client=prisma_client, - ) - - # if team change - check if this is possible - if is_different_team(data=data, existing_key_row=existing_key_row): - if llm_router is None: - raise HTTPException( - status_code=400, - detail={ - "error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI." - }, - ) - # team_obj should be set since is_different_team() returns True only when data.team_id is not None - if team_obj is None: - raise HTTPException( - status_code=500, - detail={ - "error": "Team object not found for team change validation" - }, - ) - await validate_key_team_change( - key=existing_key_row, - team=team_obj, - change_initiated_by=user_api_key_dict, - llm_router=llm_router, - ) - - # Set Management Endpoint Metadata Fields - - # Validate MCP servers in object_permission against the effective team - if data.object_permission is not None: - await _validate_mcp_servers_for_key_update( - data=data, - team_obj=team_obj, - existing_key_row=existing_key_row, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - non_default_values = await prepare_key_update_data( data=data, existing_key_row=existing_key_row )