scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

perf: batch target membership lookup in delete_user to avoid N+1

Browse source 
Greptile P1 on the /user/delete fix. Per CLAUDE.md 'No N+1 queries',
move the find_many inside the per-user loop to a single batched
fetch with {'user_id': {'in': data.user_ids}} before the loop, then
distribute to a per-user set in memory.
This commit is contained in:
user 2026-04-17 00:17:57 +00:00
parent 467166fdd7
commit 8c0668f105
No known key found for this signature in database
2 changed files with 28 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -2089,6 +2089,22 @@ async def delete_user(
},
)
# Batch-fetch target memberships once before the per-user loop. Avoids
# an N+1 DB call when delete_user is called with a large user_ids list.
target_org_ids_by_user: Dict[str, set] = {}
if not caller_is_proxy_admin:
all_target_memberships = (
await prisma_client.db.litellm_organizationmembership.find_many(
where={"user_id": {"in": data.user_ids}}
)
)
for m in all_target_memberships:
if not m.organization_id:
continue
target_org_ids_by_user.setdefault(m.user_id, set()).add(
m.organization_id
)
# check that all teams passed exist
for user_id in data.user_ids:
user_row = await prisma_client.db.litellm_usertable.find_unique(
@ -2102,14 +2118,7 @@ async def delete_user(
)
if not caller_is_proxy_admin:
target_memberships = (
await prisma_client.db.litellm_organizationmembership.find_many(
where={"user_id": user_id}
)
)
target_org_ids = {
m.organization_id for m in target_memberships if m.organization_id
}
target_org_ids = target_org_ids_by_user.get(user_id, set())
# Org-admin may only delete users whose entire org membership is
# within their admin scope. A target with ANY org outside the
# caller's scope (or no org at all) requires PROXY_ADMIN.

15
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -1919,11 +1919,18 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker):
async def mock_find_memberships(*args, **kwargs):
where = kwargs.get("where") or (args[0] if args else {})
user_id = where.get("user_id")
if user_id == "org_admin_user":
user_id_filter = where.get("user_id")
# Batched lookup: {"user_id": {"in": [...]}} returns target memberships.
# Caller role lookup: {"user_id": "<caller>", "user_role": ...}.
if isinstance(user_id_filter, dict) and "in" in user_id_filter:
if "victim" in user_id_filter["in"]:
# Attach user_id on the mock so the caller can build its
# per-user dict from the batch result.
target_membership.user_id = "victim"
return [target_membership]
return []
if user_id_filter == "org_admin_user":
return [caller_membership]
if user_id == "victim":
return [target_membership]
return []
mock_prisma_client.db.litellm_organizationmembership.find_many = mocker.AsyncMock(