From 8c0668f105c00772d785b7981973709edc718c05 Mon Sep 17 00:00:00 2001
From: user <70670632+stuxf@users.noreply.github.com>
Date: Fri, 17 Apr 2026 00:17:57 +0000
Subject: [PATCH] perf: batch target membership lookup in delete_user to avoid
 N+1

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.
---
 .../internal_user_endpoints.py                | 25 +++++++++++++------
 .../test_internal_user_endpoints.py           | 15 ++++++++---
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py
index be0f02439bc..09cd7fc5f6d 100644
--- a/litellm/proxy/management_endpoints/internal_user_endpoints.py
+++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py
@@ -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.
diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
index 104071c3e61..a074a4a211f 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
@@ -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(