diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 80094c9abd0..d30c5e62eb3 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -721,38 +721,31 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): - get all teams in LiteLLM_TeamTable - get all keys in LiteLLM_VerificationToken table - - Why separate helper for proxy admin ? - - To get Faster UI load times, get all teams and virtual keys in 1 query """ + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy.proxy_server import prisma_client - sql_query = """ - SELECT - (SELECT json_agg(t.*) FROM "LiteLLM_TeamTable" t) as teams, - (SELECT json_agg(k.*) FROM "LiteLLM_VerificationToken" k WHERE k.team_id != 'litellm-dashboard' OR k.team_id IS NULL) as keys - """ if prisma_client is None: raise Exception( "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - results = await prisma_client.db.query_raw(sql_query) - - verbose_proxy_logger.debug("results_keys: %s", results) + # Fetch teams and keys concurrently using Prisma instead of raw SQL. + _teams_in_db, keys_in_db = await asyncio.gather( + prisma_client.db.litellm_teamtable.find_many(), + prisma_client.db.litellm_verificationtoken.find_many( + where={ + "OR": [ + {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}, + {"team_id": None}, + ] + } + ), + ) - _keys_in_db: List = results[0]["keys"] or [] - # cast all keys to LiteLLM_VerificationToken - keys_in_db = [] - for key in _keys_in_db: - if key.get("models") is None: - key["models"] = [] - keys_in_db.append(LiteLLM_VerificationToken(**key)) + verbose_proxy_logger.debug("results_keys: %s", keys_in_db) - # cast all teams to LiteLLM_TeamTable - _teams_in_db: List = results[0]["teams"] or [] - _teams_in_db = [LiteLLM_TeamTable(**team) for team in _teams_in_db] _teams_in_db.sort(key=lambda x: (getattr(x, "team_alias", "") or "")) returned_keys = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db) 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 51450fd7e8b..8f31fd97c12 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 @@ -1728,4 +1728,85 @@ async def mock_find_unique(*args, **kwargs): # Verify each condition uses {"in": ["admin-creator"]} for condition in or_conditions: field = list(condition.keys())[0] - assert condition[field] == {"in": ["admin-creator"]} \ No newline at end of file + assert condition[field] == {"in": ["admin-creator"]} + + +@pytest.mark.asyncio +async def test_get_user_info_for_proxy_admin_uses_prisma(mocker): + """ + Test that _get_user_info_for_proxy_admin uses Prisma find_many calls + instead of raw SQL, and runs them concurrently. + """ + from unittest.mock import MagicMock + + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy._types import ( + LiteLLM_TeamTable, + LiteLLM_VerificationToken, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _get_user_info_for_proxy_admin, + ) + + mock_team = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One") + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.token = "hashed_token_123" + mock_key.team_id = "team-1" + mock_key.model_dump = MagicMock(return_value={ + "token": "hashed_token_123", + "team_id": "team-1", + "models": [], + }) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock( + return_value=[mock_team] + ) + mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock( + return_value=[mock_key] + ) + mock_prisma_client.get_data = mocker.AsyncMock(return_value=None) + + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ) + mocker.patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ) + mocker.patch( + "litellm.proxy.proxy_server.litellm_master_key_hash", + "master_key_hash", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user", + ) + + result = await _get_user_info_for_proxy_admin(user_api_key_dict) + + # Verify Prisma calls (not raw SQL) + mock_prisma_client.db.litellm_teamtable.find_many.assert_called_once_with() + mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with( + where={ + "OR": [ + {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}, + {"team_id": None}, + ] + } + ) + + # Verify no raw SQL was used + mock_prisma_client.db.query_raw.assert_not_called() + + # Verify response shape + assert result.user_id == "admin_user" + assert len(result.teams) == 1 + assert result.teams[0].team_id == "team-1" + assert len(result.keys) == 1 + assert result.keys[0]["team_id"] == "team-1" + assert result.keys[0]["team_alias"] == "Team One" \ No newline at end of file