From 880ee42c427ff91fed673887ceb146c9909b2956 Mon Sep 17 00:00:00 2001 From: lengkejun Date: Tue, 19 May 2026 14:12:20 +0800 Subject: [PATCH 1/2] fix(proxy): avoid N+1 key count queries in user list --- .../internal_user_endpoints.py | 34 +++++++------ .../test_internal_user_endpoints.py | 50 +++++++++++++++++++ 2 files changed, 70 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index c2f5b8eeb8b..61c59b99538 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1894,7 +1894,7 @@ async def get_user_key_counts( user_ids: list[str] | None = None, ) -> Mapping[str, int]: """ - Helper function to get the count of keys for each user using Prisma's count method. + Helper function to get the count of keys for each user using Prisma's group_by method. Args: prisma_client: The Prisma client instance @@ -1908,20 +1908,26 @@ async def get_user_key_counts( if not user_ids or len(user_ids) == 0: return {} - result: Final[dict[str, int]] = {} + result = {user_id: 0 for user_id in user_ids} - # Get count for each user_id individually - for user_id in user_ids: - count = await _verification_token_table(prisma_client).count( - where={ - "user_id": user_id, - "OR": [ - {"team_id": None}, - {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}, - ], - } - ) - result[user_id] = count + grouped_counts = await prisma_client.db.litellm_verificationtoken.group_by( + by=["user_id"], + where={ + "user_id": {"in": user_ids}, + "OR": [ + {"team_id": None}, + {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}, + ], + }, + count={"user_id": True}, + ) + + for row in grouped_counts: + user_id = row.get("user_id") + if user_id is None: + continue + count_block = row.get("_count") or {} + result[user_id] = count_block.get("user_id", 0) return result 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 f86e17c61b0..e79a69ed5a6 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 @@ -659,6 +659,56 @@ async def test_get_users_redacts_scim_enterprise_metadata(mocker): assert "scim_enterprise" not in (listed.metadata or {}) +@pytest.mark.asyncio +async def test_get_user_key_counts_returns_empty_for_no_user_ids(mocker): + mock_prisma_client = mocker.MagicMock() + + assert await get_user_key_counts(mock_prisma_client, None) == {} + assert await get_user_key_counts(mock_prisma_client, []) == {} + assert mock_prisma_client.db.litellm_verificationtoken.group_by.call_count == 0 + + +@pytest.mark.asyncio +async def test_get_user_key_counts_uses_group_by(mocker): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + mock_prisma_client = mocker.MagicMock() + captured_kwargs = {} + + async def mock_group_by(**kwargs): + captured_kwargs.update(kwargs) + return [ + {"user_id": "user-a", "_count": {"user_id": 2}}, + {"user_id": "user-c", "_count": {"user_id": 1}}, + {"user_id": "user-d"}, + {"user_id": None, "_count": {"user_id": 99}}, + ] + + async def mock_count(*args, **kwargs): + raise AssertionError("count() should not be called for user key counts") + + mock_prisma_client.db.litellm_verificationtoken.group_by = mock_group_by + mock_prisma_client.db.litellm_verificationtoken.count = mock_count + + result = await get_user_key_counts( + mock_prisma_client, + ["user-a", "user-b", "user-c", "user-d"], + ) + + assert result == {"user-a": 2, "user-b": 0, "user-c": 1, "user-d": 0} + assert captured_kwargs == { + "by": ["user_id"], + "where": { + "user_id": {"in": ["user-a", "user-b", "user-c", "user-d"]}, + "OR": [ + {"team_id": None}, + {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}, + ], + }, + "count": {"user_id": True}, + } + + def test_validate_sort_params(): """ Test that validate_sort_params returns None if sort_by is None From eba2685ece06e82b2283f435fbeb1afd4d2ccd81 Mon Sep 17 00:00:00 2001 From: lengkejun Date: Thu, 21 May 2026 14:55:09 +0800 Subject: [PATCH 2/2] docs(proxy): refresh get_user_key_counts docstring after group_by switch --- .../proxy/management_endpoints/internal_user_endpoints.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 61c59b99538..e1d7c751016 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1894,14 +1894,16 @@ async def get_user_key_counts( user_ids: list[str] | None = None, ) -> Mapping[str, int]: """ - Helper function to get the count of keys for each user using Prisma's group_by method. + Get the count of (non-UI-session) keys for each user via a single + Prisma ``group_by`` query, avoiding an N+1 per-user ``count`` loop. Args: prisma_client: The Prisma client instance user_ids: List of user IDs to get key counts for Returns: - Dictionary mapping user_id to key count + Dictionary mapping user_id to key count. Users in ``user_ids`` with no + keys are present with a value of ``0``. """ from litellm.constants import UI_SESSION_TOKEN_TEAM_ID