use raw SQL in key_aliases to avoid loading full rows into memory

Replace Prisma ORM count/find_many calls with two query_raw calls that
only project the key_alias column. The Prisma client wrapper does not
support SELECT projection via find_many, so raw SQL is used to keep
memory usage proportional to the page size rather than total key count.
Update tests to mock query_raw instead of count/find_many.

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>

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(

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