Avoid redundant bulk team key lookups

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -1889,6 +1889,7 @@ async def _process_single_key_update(
     proxy_logging_obj: Any,
     llm_router: Optional[Router],
     user_custom_key_update: Optional[Callable] = None,
+    existing_key_row: Optional[LiteLLM_VerificationToken] = None,
 ) -> Dict[str, Any]:
     """
     Process a single key update with all validations and checks.
@@ -1904,6 +1905,7 @@ async def _process_single_key_update(
         user_api_key_cache: User API key cache
         proxy_logging_obj: Proxy logging object
         llm_router: LLM router instance
+        existing_key_row: Optional pre-fetched key row to avoid redundant lookups
 
     Returns:
         Dict containing the updated key information
@@ -1915,10 +1917,11 @@ async def _process_single_key_update(
     _validate_max_budget(update_key_request.max_budget)
 
     # Get and validate existing key
-    existing_key_row = await _get_and_validate_existing_key(
-        token=update_key_request.key,
-        prisma_client=prisma_client,
-    )
+    if existing_key_row is None:
+        existing_key_row = await _get_and_validate_existing_key(
+            token=update_key_request.key,
+            prisma_client=prisma_client,
+        )
 
     # Check team member permissions
     if prisma_client is not None:
@@ -2852,6 +2855,7 @@ async def bulk_update_team_keys(
                 proxy_logging_obj=proxy_logging_obj,
                 llm_router=llm_router,
                 user_custom_key_update=user_custom_key_update,
+                existing_key_row=existing_by_token[db_token],
             )
 
             successful_updates.append(

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -10117,10 +10117,11 @@ async def test_bulk_update_team_keys_success_with_key_ids(monkeypatch):
     )
 
     keys = [_make_team_key("tok-a"), _make_team_key("tok-b")]
+    find_unique = AsyncMock(side_effect=keys)
     mock = _setup_team_keys_mocks(
         monkeypatch,
         find_many=keys,
-        find_unique=AsyncMock(side_effect=keys),
+        find_unique=find_unique,
         update_data=AsyncMock(
             side_effect=[{"data": _updated({"max_budget": 50.0})}] * 2
         ),
@@ -10139,6 +10140,7 @@ async def test_bulk_update_team_keys_success_with_key_ids(monkeypatch):
     where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
     assert where["team_id"] == "team-abc"
     assert where["token"] == {"in": ["tok-a", "tok-b"]}
+    find_unique.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -10149,10 +10151,11 @@ async def test_bulk_update_team_keys_success_all_keys_in_team(monkeypatch):
     )
 
     keys = [_make_team_key(f"tok-{i}") for i in range(3)]
+    find_unique = AsyncMock(side_effect=keys)
     mock = _setup_team_keys_mocks(
         monkeypatch,
         find_many=keys,
-        find_unique=AsyncMock(side_effect=keys),
+        find_unique=find_unique,
         update_data=AsyncMock(
             side_effect=[{"data": _updated({"max_budget": 50.0})}] * 3
         ),
@@ -10178,6 +10181,7 @@ async def test_bulk_update_team_keys_success_all_keys_in_team(monkeypatch):
         for c in expires_or
         if isinstance(c.get("expires"), dict)
     )
+    find_unique.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -10479,8 +10483,7 @@ async def test_bulk_update_team_keys_does_not_log_raw_sk_token_on_failure(
     _setup_team_keys_mocks(
         monkeypatch,
         find_many=[row],
-        # Force the per-key fetch to raise so the exception logger fires.
-        find_unique=AsyncMock(side_effect=RuntimeError("boom")),
+        update_data=AsyncMock(side_effect=RuntimeError("boom")),
         hash_identity=False,
     )