diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8a5145bf793..a0442a476f9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4135,29 +4135,41 @@ async def key_aliases( verbose_proxy_logger.error("Database not connected") raise Exception("Database not connected") - conditions: List[Dict[str, Any]] = [{"key_alias": {"not": None}}] - try: - conditions.append(_get_condition_to_filter_out_ui_session_tokens()) - except NameError: - # Helper may not exist in some builds; ignore if missing - pass + # Build a parameterized WHERE clause to avoid loading full rows into + # memory. Raw SQL is used because the Prisma client wrapper does not + # support column-level SELECT projection on find_many. + # + # $1 is always UI_SESSION_TOKEN_TEAM_ID (filters out UI session tokens). + query_params: List[Any] = [UI_SESSION_TOKEN_TEAM_ID] + where_parts = [ + "key_alias IS NOT NULL", + "key_alias != ''", + "(team_id IS NULL OR team_id != $1)", + ] if search: - conditions.append( - {"key_alias": {"contains": search, "mode": "insensitive"}} - ) - where: Dict[str, Any] = {"AND": conditions} + query_params.append(f"%{search}%") + where_parts.append(f"key_alias ILIKE ${len(query_params)}") - total_count = await prisma_client.db.litellm_verificationtoken.count( - where=where, - ) - rows = await prisma_client.db.litellm_verificationtoken.find_many( - where=where, - order=[{"key_alias": "asc"}], - skip=(page - 1) * size, - take=size, - ) + where_sql = " AND ".join(where_parts) - aliases: List[str] = [row.key_alias for row in rows if row.key_alias] # type: ignore[misc] + count_sql = ( + f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' + ) + count_rows = await prisma_client.db.query_raw(count_sql, *query_params) + total_count = int(count_rows[0]["count"]) if count_rows else 0 + + aliases_params = query_params + [size, (page - 1) * size] + limit_idx = len(aliases_params) - 1 + offset_idx = len(aliases_params) + aliases_sql = ( + f"SELECT key_alias" + f' FROM "LiteLLM_VerificationToken"' + f" WHERE {where_sql}" + f" ORDER BY key_alias ASC" + f" LIMIT ${limit_idx} OFFSET ${offset_idx}" + ) + alias_rows = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + aliases: List[str] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages = -(-total_count // size) if total_count > 0 else 0 verbose_proxy_logger.debug( diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index bdfc6734664..75325508b01 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6217,15 +6217,12 @@ async def test_generate_key_helper_fn_agent_id(): @pytest.mark.asyncio async def test_key_aliases_response_shape(): """Test that key_aliases returns the correct paginated response shape.""" - mock_row1 = MagicMock() - mock_row1.key_alias = "alias-alpha" - mock_row2 = MagicMock() - mock_row2.key_alias = "alias-beta" - mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=2) - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[mock_row1, mock_row2] + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 2}], + [{"key_alias": "alias-alpha"}, {"key_alias": "alias-beta"}], + ] ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): @@ -6237,21 +6234,23 @@ async def test_key_aliases_response_shape(): assert result["total_pages"] == 1 assert result["size"] == 50 - # Both count and find_many must use the same where clause - count_where = mock_prisma_client.db.litellm_verificationtoken.count.call_args.kwargs["where"] - find_where = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs["where"] - assert count_where == find_where - - # Non-null alias filter must be present - assert json.dumps({"key_alias": {"not": None}}) in json.dumps(count_where) + # Both SQL calls must filter out null/empty aliases + count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] + aliases_sql = mock_prisma_client.db.query_raw.call_args_list[1].args[0] + assert "key_alias IS NOT NULL" in count_sql + assert "key_alias IS NOT NULL" in aliases_sql @pytest.mark.asyncio async def test_key_aliases_pagination_skip_take(): - """Test that skip and take are correctly computed from page and size.""" + """Test that LIMIT and OFFSET are correctly derived from page and size.""" mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=120) - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 120}], + [], + ] + ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): result = await key_aliases(page=3, size=25, search=None) @@ -6261,39 +6260,49 @@ async def test_key_aliases_pagination_skip_take(): assert result["total_count"] == 120 assert result["total_pages"] == 5 # ceil(120 / 25) - find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs - assert find_many_kwargs["skip"] == 50 # (3 - 1) * 25 - assert find_many_kwargs["take"] == 25 + # aliases query params: [UI_SESSION_TOKEN_TEAM_ID, size=25, offset=50] + aliases_call_args = mock_prisma_client.db.query_raw.call_args_list[1].args + assert aliases_call_args[-2] == 25 # LIMIT = size + assert aliases_call_args[-1] == 50 # OFFSET = (3 - 1) * 25 @pytest.mark.asyncio async def test_key_aliases_search_filter(): - """Test that the search param adds a case-insensitive contains condition.""" + """Test that the search param adds a case-insensitive ILIKE condition.""" mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 0}], + [], + ] + ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): await key_aliases(page=1, size=50, search="my-key") - where = mock_prisma_client.db.litellm_verificationtoken.count.call_args.kwargs["where"] - assert ( - json.dumps({"key_alias": {"contains": "my-key", "mode": "insensitive"}}) - in json.dumps(where) - ) + count_call = mock_prisma_client.db.query_raw.call_args_list[0] + count_sql = count_call.args[0] + count_params = count_call.args[1:] + + assert "ILIKE" in count_sql + assert "%my-key%" in count_params @pytest.mark.asyncio -async def test_key_aliases_no_search_omits_contains_filter(): - """Test that without a search term no contains condition is added.""" +async def test_key_aliases_no_search_omits_ilike_filter(): + """Test that without a search term no ILIKE condition is added.""" mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 0}], + [], + ] + ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): await key_aliases(page=1, size=50, search=None) - where = mock_prisma_client.db.litellm_verificationtoken.count.call_args.kwargs["where"] - assert "contains" not in json.dumps(where) + count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] + assert "ILIKE" not in count_sql