mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(proxy): filter /key/list by user email via the user relation
Adds a user_email query param resolved as a Prisma relation condition, so Postgres semi-joins keys to users in one query with complete results and nothing materialized in the proxy. The deleted-keys table has no user relation (its rows may reference vanished users), so that path resolves ids with a bounded raw lookup that mirrors the relation filter's wildcard semantics. The Admin UI Virtual Keys filter drawer gains a User Email field wired through useKeys
This commit is contained in:
parent
af05f4b249
commit
055f8f067d
5 changed files with 290 additions and 3 deletions
|
|
@ -5199,6 +5199,34 @@ async def get_member_team_ids(
|
|||
|
||||
|
||||
VALID_EXPIRES_FILTER_VALUES = frozenset({"active", "expired"})
|
||||
USER_EMAIL_KEY_FILTER_MAX_USERS = 50_000
|
||||
|
||||
|
||||
async def _resolve_user_ids_for_email_filter(prisma_client: PrismaClient, user_email: str) -> list[str]:
|
||||
"""Resolve a user_email fragment to the complete set of matching user IDs.
|
||||
|
||||
Only used for the deleted-keys table, which has no relation to the user
|
||||
table; live keys filter via the litellm_user_table relation instead. Raw
|
||||
SQL keeps this bounded: only the user_id column is fetched (find_many
|
||||
materializes full rows). The fragment is passed to ILIKE verbatim so LIKE
|
||||
wildcards behave exactly like the live path and the Users page email
|
||||
search, both of which use Prisma contains, which passes wildcards through.
|
||||
Matching more than USER_EMAIL_KEY_FILTER_MAX_USERS users is refused rather
|
||||
than truncated: a silently partial id list would make /key/list drop
|
||||
visible keys.
|
||||
"""
|
||||
rows = await prisma_client.db.query_raw(
|
||||
'SELECT user_id FROM "LiteLLM_UserTable" WHERE user_email ILIKE $1 LIMIT $2',
|
||||
f"%{user_email}%",
|
||||
USER_EMAIL_KEY_FILTER_MAX_USERS + 1,
|
||||
)
|
||||
user_ids = [row["user_id"] for row in rows]
|
||||
if len(user_ids) > USER_EMAIL_KEY_FILTER_MAX_USERS:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "user_email filter matched too many users; use a more specific email fragment"},
|
||||
)
|
||||
return user_ids
|
||||
|
||||
|
||||
@router.get(
|
||||
|
|
@ -5216,6 +5244,10 @@ async def list_keys(
|
|||
None,
|
||||
description="Filter keys by user ID. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.",
|
||||
),
|
||||
user_email: str | None = Query(
|
||||
None,
|
||||
description="Filter keys by the owning user's email. Case-insensitive substring match against the user table; only keys whose user_id belongs to a matching user are returned. For status=deleted, fragments matching more than 50,000 users are rejected with a 400.",
|
||||
),
|
||||
team_id: Optional[str] = Query(None, description="Filter keys by team ID"),
|
||||
organization_id: Optional[str] = Query(None, description="Filter keys by organization ID"),
|
||||
key_hash: Optional[str] = Query(None, description="Filter keys by key hash"),
|
||||
|
|
@ -5366,6 +5398,7 @@ async def list_keys(
|
|||
agent_id=agent_id,
|
||||
use_substring_matching=use_substring_matching,
|
||||
expires_filter=expires if isinstance(expires, str) else None,
|
||||
user_email_filter=user_email if user_email and isinstance(user_email, str) else None,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Successfully prepared response")
|
||||
|
|
@ -5594,6 +5627,8 @@ def _build_key_filter_conditions(
|
|||
agent_id: Optional[str] = None,
|
||||
use_substring_matching: bool = False,
|
||||
expires_filter: str | None = None,
|
||||
user_ids_for_email: list[str] | None = None,
|
||||
user_email_filter: str | None = None,
|
||||
) -> Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]]:
|
||||
"""Build filter conditions for key listing.
|
||||
|
||||
|
|
@ -5701,6 +5736,15 @@ def _build_key_filter_conditions(
|
|||
where = {"AND": [where, {"access_group_ids": {"hasSome": [access_group_id]}}]}
|
||||
if agent_id and isinstance(agent_id, str):
|
||||
where = {"AND": [where, {"agent_id": agent_id}]}
|
||||
if user_ids_for_email is not None:
|
||||
where = {"AND": [where, {"user_id": {"in": user_ids_for_email}}]}
|
||||
if user_email_filter:
|
||||
where = {
|
||||
"AND": [
|
||||
where,
|
||||
{"litellm_user_table": {"is": {"user_email": {"contains": user_email_filter, "mode": "insensitive"}}}},
|
||||
]
|
||||
}
|
||||
if expires_filter is not None and expires_filter in VALID_EXPIRES_FILTER_VALUES:
|
||||
where = {"AND": [where, _build_expires_where_clause(expires_filter, datetime.now(timezone.utc))]}
|
||||
|
||||
|
|
@ -5733,6 +5777,7 @@ async def _list_key_helper(
|
|||
agent_id: Optional[str] = None,
|
||||
use_substring_matching: bool = False,
|
||||
expires_filter: str | None = None,
|
||||
user_email_filter: str | None = None,
|
||||
) -> KeyListResponseObject:
|
||||
"""
|
||||
Helper function to list keys
|
||||
|
|
@ -5756,6 +5801,18 @@ async def _list_key_helper(
|
|||
"total_pages": int,
|
||||
}
|
||||
"""
|
||||
# Determine which table to query based on status
|
||||
use_deleted_table = status == "deleted"
|
||||
|
||||
# The deleted-keys table has no relation to the user table (its rows may
|
||||
# reference users that no longer exist), so an email filter there falls
|
||||
# back to resolving user ids up front; live keys filter via the relation.
|
||||
user_ids_for_email = (
|
||||
await _resolve_user_ids_for_email_filter(prisma_client, user_email_filter)
|
||||
if user_email_filter is not None and use_deleted_table
|
||||
else None
|
||||
)
|
||||
|
||||
where = _build_key_filter_conditions(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
|
|
@ -5771,6 +5828,8 @@ async def _list_key_helper(
|
|||
agent_id=agent_id,
|
||||
use_substring_matching=use_substring_matching,
|
||||
expires_filter=expires_filter,
|
||||
user_ids_for_email=user_ids_for_email,
|
||||
user_email_filter=user_email_filter if not use_deleted_table else None,
|
||||
)
|
||||
|
||||
# Calculate skip for pagination
|
||||
|
|
@ -5782,9 +5841,6 @@ async def _list_key_helper(
|
|||
_validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None
|
||||
)
|
||||
|
||||
# Determine which table to query based on status
|
||||
use_deleted_table = status == "deleted"
|
||||
|
||||
# Fetch keys with pagination
|
||||
if use_deleted_table:
|
||||
keys = await DeletedVerificationTokenRepository(prisma_client).table.find_many(
|
||||
|
|
|
|||
|
|
@ -6257,6 +6257,224 @@ def test_build_key_filter_conditions_agent_id_narrows_visibility():
|
|||
assert "agent_id" not in json.dumps(where_without)
|
||||
|
||||
|
||||
def test_build_key_filter_conditions_user_email_ids_narrow_visibility():
|
||||
"""
|
||||
Deleted-key listings resolve user_email to a list of user IDs which must
|
||||
be ANDed on top of the caller's visibility conditions (narrow, never widen).
|
||||
An email that matches no users must produce a match-nothing filter rather
|
||||
than being dropped (dropping it would return every visible key).
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_build_key_filter_conditions,
|
||||
)
|
||||
|
||||
where = _build_key_filter_conditions(
|
||||
user_id="some-user",
|
||||
team_id=None,
|
||||
organization_id=None,
|
||||
key_alias=None,
|
||||
key_hash=None,
|
||||
exclude_team_id=None,
|
||||
admin_team_ids=["team-a"],
|
||||
member_team_ids=None,
|
||||
include_created_by_keys=False,
|
||||
user_ids_for_email=["user-1", "user-2"],
|
||||
)
|
||||
assert where.get("AND"), f"expected top-level AND, got: {where}"
|
||||
assert {"user_id": {"in": ["user-1", "user-2"]}} in where["AND"], f"user_id in-filter not ANDed: {where}"
|
||||
|
||||
where_no_match = _build_key_filter_conditions(
|
||||
user_id=None,
|
||||
team_id=None,
|
||||
organization_id=None,
|
||||
key_alias=None,
|
||||
key_hash=None,
|
||||
exclude_team_id=None,
|
||||
admin_team_ids=None,
|
||||
member_team_ids=None,
|
||||
include_created_by_keys=False,
|
||||
user_ids_for_email=[],
|
||||
)
|
||||
assert {"user_id": {"in": []}} in where_no_match["AND"], (
|
||||
f"an email matching no users must match nothing, got: {where_no_match}"
|
||||
)
|
||||
|
||||
where_without = _build_key_filter_conditions(
|
||||
user_id="some-user",
|
||||
team_id=None,
|
||||
organization_id=None,
|
||||
key_alias=None,
|
||||
key_hash=None,
|
||||
exclude_team_id=None,
|
||||
admin_team_ids=None,
|
||||
member_team_ids=None,
|
||||
include_created_by_keys=False,
|
||||
)
|
||||
assert '"in"' not in json.dumps(where_without)
|
||||
assert "litellm_user_table" not in json.dumps(where_without)
|
||||
|
||||
|
||||
def test_build_key_filter_conditions_user_email_relation_filter():
|
||||
"""
|
||||
A live-table user_email filter must be expressed as a relation condition
|
||||
(DB-side semi-join to the user table) ANDed on top of the caller's
|
||||
visibility conditions, so results are complete without materializing
|
||||
user ids in the proxy.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_build_key_filter_conditions,
|
||||
)
|
||||
|
||||
where = _build_key_filter_conditions(
|
||||
user_id="some-user",
|
||||
team_id=None,
|
||||
organization_id=None,
|
||||
key_alias=None,
|
||||
key_hash=None,
|
||||
exclude_team_id=None,
|
||||
admin_team_ids=["team-a"],
|
||||
member_team_ids=None,
|
||||
include_created_by_keys=False,
|
||||
user_email_filter="alias@example.com",
|
||||
)
|
||||
assert where.get("AND"), f"expected top-level AND, got: {where}"
|
||||
assert {
|
||||
"litellm_user_table": {"is": {"user_email": {"contains": "alias@example.com", "mode": "insensitive"}}}
|
||||
} in where["AND"], f"relation filter not ANDed: {where}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_keys_user_email_passes_filter_to_helper():
|
||||
"""
|
||||
list_keys must pass the user_email filter through to _list_key_helper
|
||||
verbatim without resolving anything itself. No email filter (including
|
||||
the Query default object from direct calls) must pass None.
|
||||
"""
|
||||
from unittest.mock import Mock
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_query_raw = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.query_raw = mock_query_raw
|
||||
mock_list_key_helper = AsyncMock(
|
||||
return_value={"keys": [], "total_count": 0, "current_page": 1, "total_pages": 0}
|
||||
)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.validate_key_list_check",
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper",
|
||||
mock_list_key_helper,
|
||||
),
|
||||
):
|
||||
admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
|
||||
|
||||
await list_keys(
|
||||
request=Mock(),
|
||||
user_api_key_dict=admin,
|
||||
user_email="alias@example.com",
|
||||
include_team_keys=False,
|
||||
include_created_by_keys=False,
|
||||
status=None,
|
||||
)
|
||||
mock_query_raw.assert_not_called()
|
||||
assert mock_list_key_helper.call_args.kwargs["user_email_filter"] == "alias@example.com"
|
||||
|
||||
await list_keys(
|
||||
request=Mock(),
|
||||
user_api_key_dict=admin,
|
||||
include_team_keys=False,
|
||||
include_created_by_keys=False,
|
||||
status=None,
|
||||
)
|
||||
assert mock_list_key_helper.call_args.kwargs["user_email_filter"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_key_helper_user_email_routes_live_vs_deleted():
|
||||
"""
|
||||
Live listings must filter via the user relation (no id materialization);
|
||||
deleted listings have no relation, so the email must resolve to user ids
|
||||
first and land as a user_id in-filter.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_query_raw = AsyncMock(return_value=[{"user_id": "user-1"}])
|
||||
mock_prisma_client.db.query_raw = mock_query_raw
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0)
|
||||
mock_prisma_client.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_deletedverificationtoken.count = AsyncMock(return_value=0)
|
||||
|
||||
await _list_key_helper(
|
||||
prisma_client=mock_prisma_client,
|
||||
page=1,
|
||||
size=10,
|
||||
user_id=None,
|
||||
team_id=None,
|
||||
organization_id=None,
|
||||
key_alias=None,
|
||||
key_hash=None,
|
||||
user_email_filter="alias@example.com",
|
||||
)
|
||||
mock_query_raw.assert_not_called()
|
||||
live_where = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs["where"]
|
||||
assert {
|
||||
"litellm_user_table": {"is": {"user_email": {"contains": "alias@example.com", "mode": "insensitive"}}}
|
||||
} in live_where["AND"], f"live listing must use the relation filter: {live_where}"
|
||||
|
||||
await _list_key_helper(
|
||||
prisma_client=mock_prisma_client,
|
||||
page=1,
|
||||
size=10,
|
||||
user_id=None,
|
||||
team_id=None,
|
||||
organization_id=None,
|
||||
key_alias=None,
|
||||
key_hash=None,
|
||||
status="deleted",
|
||||
user_email_filter="alias@example.com",
|
||||
)
|
||||
mock_query_raw.assert_called_once()
|
||||
deleted_where = mock_prisma_client.db.litellm_deletedverificationtoken.find_many.call_args.kwargs["where"]
|
||||
assert {"user_id": {"in": ["user-1"]}} in deleted_where["AND"], (
|
||||
f"deleted listing must fall back to the resolved id in-filter: {deleted_where}"
|
||||
)
|
||||
assert "litellm_user_table" not in json.dumps(deleted_where), (
|
||||
"deleted table has no user relation; a relation filter there would error"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_user_ids_for_email_filter_wildcard_parity_and_refuses_overflow():
|
||||
"""
|
||||
The fragment must reach ILIKE verbatim so wildcard behavior matches the
|
||||
live relation path and the Users page email search (Prisma contains
|
||||
passes LIKE wildcards through), and a fragment matching more than the
|
||||
ceiling must 400 instead of silently truncating the id list (truncation
|
||||
would make /key/list drop visible keys).
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
USER_EMAIL_KEY_FILTER_MAX_USERS,
|
||||
_resolve_user_ids_for_email_filter,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_query_raw = AsyncMock(return_value=[{"user_id": "user-1"}])
|
||||
mock_prisma_client.db.query_raw = mock_query_raw
|
||||
|
||||
result = await _resolve_user_ids_for_email_filter(mock_prisma_client, "50%_off\\weird")
|
||||
assert result == ["user-1"]
|
||||
assert mock_query_raw.call_args.args[1] == "%50%_off\\weird%"
|
||||
|
||||
mock_query_raw.return_value = [{"user_id": f"user-{i}"} for i in range(USER_EMAIL_KEY_FILTER_MAX_USERS + 1)]
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _resolve_user_ids_for_email_filter(mock_prisma_client, "@")
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "too many users" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_user_row_for_key_write_create_only_upsert():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ export interface KeyListCallOptions {
|
|||
agentID?: string | null;
|
||||
selectedKeyAlias?: string | null;
|
||||
userID?: string | null;
|
||||
userEmail?: string | null;
|
||||
keyHash?: string | null;
|
||||
sortBy?: string | null;
|
||||
sortOrder?: string | null;
|
||||
|
|
@ -55,6 +56,7 @@ const keyListCall = async (accessToken: string, page: number, pageSize: number,
|
|||
key_alias: options.selectedKeyAlias,
|
||||
key_hash: options.keyHash,
|
||||
user_id: options.userID,
|
||||
user_email: options.userEmail,
|
||||
page,
|
||||
size: pageSize,
|
||||
sort_by: options.sortBy,
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@ const FILTER_LABELS: Record<string, string> = {
|
|||
team_id: "Team",
|
||||
org_id: "Organization",
|
||||
user_id: "User ID",
|
||||
user_email: "User Email",
|
||||
key_hash: "Key ID",
|
||||
};
|
||||
|
||||
|
|
@ -73,6 +74,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
|
|||
organizationID: getFilterValue("org_id"),
|
||||
selectedKeyAlias: searchQuery.trim() || undefined,
|
||||
userID: getFilterValue("user_id"),
|
||||
userEmail: getFilterValue("user_email"),
|
||||
keyHash: getFilterValue("key_hash"),
|
||||
sortBy,
|
||||
sortOrder,
|
||||
|
|
@ -247,6 +249,13 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
|
|||
placeholder="Enter User ID…"
|
||||
/>
|
||||
</DataTableFilterField>
|
||||
<DataTableFilterField label="User Email">
|
||||
<Input
|
||||
value={(get("user_email") as string) ?? ""}
|
||||
onChange={(event) => set("user_email", event.target.value)}
|
||||
placeholder="Enter User Email…"
|
||||
/>
|
||||
</DataTableFilterField>
|
||||
<DataTableFilterField label="Key ID">
|
||||
<Input
|
||||
value={(get("key_hash") as string) ?? ""}
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -42915,6 +42915,8 @@ export interface operations {
|
|||
size?: number;
|
||||
/** @description Filter keys by user ID. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching. */
|
||||
user_id?: string | null;
|
||||
/** @description Filter keys by the owning user's email. Case-insensitive substring match against the user table; only keys whose user_id belongs to a matching user are returned. For status=deleted, fragments matching more than 50,000 users are rejected with a 400. */
|
||||
user_email?: string | null;
|
||||
/** @description Filter keys by team ID */
|
||||
team_id?: string | null;
|
||||
/** @description Filter keys by organization ID */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue