scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

refactor(mcp): batch the purge row deletion into one query

Browse source 
The per-row delete_many loop becomes a single delete filtered to the enumerated OAuth users'
(user_id IN, server_id) pairs; same rows deleted, same BYOK-sparing precision, same count-mismatch
detection, one round-trip instead of N
This commit is contained in:
Tin 2026-07-09 16:18:25 -07:00
parent c46c9d4652
commit 9dcc21cd48
2 changed files with 9 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -1175,8 +1175,8 @@ async def purge_user_oauth_credentials_for_server(
oauth_rows = [row for row in rows if _decode_oauth_payload(row.credential_b64) is not None]
if not oauth_rows:
return 0
deleted_count = sum(
[await repo.table.delete_many(where={"user_id": row.user_id, "server_id": server_id}) for row in oauth_rows]
deleted_count = await repo.table.delete_many(
where={"server_id": server_id, "user_id": {"in": [row.user_id for row in oauth_rows]}}
)
if invalidate_token_cache is None:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (

12
tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py
View file

@ -171,7 +171,7 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_each_user():
prisma = MagicMock()
prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[_oauth_row("alice"), _oauth_row("bob")])
prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=1)
prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2)
invalidations = []
@ -181,15 +181,17 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_each_user():
purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=record_invalidation)
assert purged == 2
assert prisma.db.litellm_mcpusercredentials.delete_many.await_count == 2
prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once_with(
where={"server_id": "srv-1", "user_id": {"in": ["alice", "bob"]}}
)
assert set(invalidations) == {("alice", "srv-1"), ("bob", "srv-1")}
@pytest.mark.asyncio
async def test_purge_user_oauth_credentials_for_server_spares_byok_rows():
"""Regression: the purge used to delete_many on server_id alone, wiping BYOK API keys that share
the LiteLLM_MCPUserCredentials table. Only rows holding an OAuth2 payload may be deleted, each by
its (user_id, server_id) pair, and only their users' token caches invalidated."""
the LiteLLM_MCPUserCredentials table. Only rows holding an OAuth2 payload may be deleted (one
batched query filtered to their user_ids), and only their users' token caches invalidated."""
from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server
prisma = MagicMock()
@ -205,7 +207,7 @@ async def test_purge_user_oauth_credentials_for_server_spares_byok_rows():
assert purged == 1
prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once_with(
where={"user_id": "alice", "server_id": "srv-1"}
where={"server_id": "srv-1", "user_id": {"in": ["alice"]}}
)
assert invalidations == [("alice", "srv-1")]