diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 9a98bdbb6b1..70ddf998599 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1889,34 +1889,42 @@ 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. + 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 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 8bce967b316..e8a4321ba5e 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