diff --git a/litellm/proxy/management_helpers/bulk_user_deletion.py b/litellm/proxy/management_helpers/bulk_user_deletion.py index e7726652eca..d73901a3ea0 100644 --- a/litellm/proxy/management_helpers/bulk_user_deletion.py +++ b/litellm/proxy/management_helpers/bulk_user_deletion.py @@ -64,7 +64,7 @@ if TYPE_CHECKING: from litellm.repositories.prisma_protocols import TableActions _AUDIT_LOG_CONCURRENCY: Final = 10 -_USER_BATCH_TX_TIMEOUT: Final = timedelta(seconds=60) +_BATCH_TX_TIMEOUT: Final = timedelta(seconds=60) class _ErrorDetail(TypedDict): @@ -289,7 +289,7 @@ async def bulk_remove_team_members( duplicates: Final = _duplicate_member_indexes(data.members) kept_indexes: Final = tuple(i for i in range(len(data.members)) if i not in duplicates) members: Final = tuple(data.members[i] for i in kept_indexes) - async with prisma_client.tx() as tx: + async with prisma_client.tx(timeout=_BATCH_TX_TIMEOUT) as tx: removal: Final = await _remove_members_from_team(prisma_client, tx, data.team_id, members, user_api_key_dict) await delete_cache_key_objects( hashed_tokens=removal.deleted_key_tokens, @@ -384,7 +384,7 @@ async def _delete_users_tx( """Rewrites every team the users belong to and deletes their rows in one transaction, so a failure anywhere rolls back the whole batch. Teams a user still names but which no longer exist are skipped; the user row goes away regardless.""" - async with prisma_client.tx(timeout=_USER_BATCH_TX_TIMEOUT) as tx: + async with prisma_client.tx(timeout=_BATCH_TX_TIMEOUT) as tx: team_rows: Final = await _team_tx_db(tx).find_many( where=_in_filter("team_id", frozenset(t for teams in teams_of.values() for t in teams)) ) @@ -395,11 +395,7 @@ async def _delete_users_tx( prisma_client, tx, tid, - tuple( - MemberDeleteRequest(user_id=u.user_id, user_email=u.user_email) - for u in users - if tid in teams_of[u.user_id] - ), + tuple(MemberDeleteRequest(user_id=u.user_id) for u in users if tid in teams_of[u.user_id]), user_api_key_dict, ) for tid in team_ids diff --git a/tests/test_litellm/proxy/management_helpers/test_bulk_user_deletion.py b/tests/test_litellm/proxy/management_helpers/test_bulk_user_deletion.py index dd266880ce0..aac30184e1d 100644 --- a/tests/test_litellm/proxy/management_helpers/test_bulk_user_deletion.py +++ b/tests/test_litellm/proxy/management_helpers/test_bulk_user_deletion.py @@ -277,6 +277,37 @@ async def test_bulk_delete_removes_users_from_every_team_and_store(): assert prisma.locks == ["t1", "t2"] and prisma.roster_reads == ["t1", "t2"] +@pytest.mark.asyncio +async def test_bulk_delete_leaves_teammates_who_share_the_deleted_users_email_alone(): + twin = _UserRow(user_id="twin", user_email="u1@example.com", teams=["t1"]) + team = LiteLLM_TeamTable( + team_id="t1", + members_with_roles=[ + Member(user_id="u1", user_email="u1@example.com", role="user"), + Member(user_id="twin", user_email="u1@example.com", role="user"), + ], + ) + prisma = _FakePrisma( + users=[_user("u1", "t1"), twin], + teams=[team], + memberships=[("t1", "u1"), ("t1", "twin")], + tokens=[ + {"token": "k1", "user_id": "u1", "team_id": "t1"}, + {"token": "k-twin", "user_id": "twin", "team_id": "t1"}, + ], + ) + + response = await _delete(prisma, ["u1"]) + + assert [(r.success, r.teams_removed) for r in response.results] == [(True, ("t1",))] + assert _roster(prisma, "t1") == ["twin"] + assert set(prisma.db.litellm_usertable.rows) == {"twin"} and prisma.db.litellm_usertable.rows["twin"].teams == [ + "t1" + ] + assert prisma.db.litellm_teammembership.rows == [{"team_id": "t1", "user_id": "twin"}] + assert [t["token"] for t in prisma.db.litellm_verificationtoken.rows] == ["k-twin"] + + @pytest.mark.asyncio async def test_bulk_delete_finds_teams_through_membership_rows_when_user_teams_array_is_stale(): prisma = _FakePrisma(