diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index 3b587e932d4..8bd7e83d122 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -31,6 +31,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import invalidate_team_member_spend_state from litellm.proxy.auth.litellm_license import LicenseCheck from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks from litellm.proxy.management_endpoints.common_utils import ( _is_user_org_admin_for_team, # pyright: ignore[reportPrivateUsage] # same team-admin check /user/new uses @@ -433,23 +434,38 @@ async def _insert_users( try: await table.create_many(data=payloads) return tuple(prepared), () - except Exception: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified + except Exception as exc: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified verbose_proxy_logger.warning("/user/bulk_new: create_many failed, retrying rows individually", exc_info=True) - landed_rows: Final = await table.find_many( - where={"user_id": {"in": [payload["user_id"] for payload in payloads]}} # mutable-ok: Prisma filter - ) + outcome_unknown: Final = PrismaDBExceptionHandler.is_database_infrastructure_error(exc) + requested: Final = frozenset(payload["user_id"] for payload in payloads) + landed_rows: Final = await table.find_many(where={"user_id": {"in": list(requested)}}) # mutable-ok: Prisma filter landed: Final = frozenset(row.user_id for row in landed_rows) + # create_many is one INSERT: after a lost response the full set is ours, any partial set belongs to another request + if outcome_unknown and landed == requested: + return tuple(prepared), () + taken: Final = tuple(user for user in prepared if user.row.user_id in landed) retried: Final = tuple(user for user in prepared if user.row.user_id not in landed) outcomes: Final = await _bounded( BULK_NEW_USER_CONCURRENCY, tuple(table.create(data=_user_create_payload(user)) for user in retried) ) failed: Final = MappingProxyType( { - user.row.user_id: _RowFailure( - user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome) - ) - for user, outcome in zip(retried, outcomes, strict=True) - if isinstance(outcome, BaseException) + **{ + user.row.user_id: _RowFailure( + user.pending.index, + user.pending.user_id, + user.row.user_email, + f"User id={user.row.user_id} already exists", + ) + for user in taken + }, + **{ + user.row.user_id: _RowFailure( + user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome) + ) + for user, outcome in zip(retried, outcomes, strict=True) + if isinstance(outcome, BaseException) + }, } ) return ( diff --git a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py index 7ff9b037cc1..00ddd63c057 100644 --- a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py @@ -2,8 +2,10 @@ import json from contextlib import asynccontextmanager from typing import Final +import httpx import pytest from fastapi import HTTPException +from prisma.errors import UniqueViolationError from pydantic import BaseModel, ConfigDict, ValidationError from litellm.caching.caching import DualCache @@ -31,10 +33,16 @@ class _UserRow(BaseModel): class _UserTable: """Enough of the Prisma user table for the bulk path: set lookups, one create_many and per-row fallbacks.""" - def __init__(self, fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False) -> None: + def __init__( + self, + fail_ids: frozenset[str] = frozenset(), + commit_then_drop: bool = False, + raced_ids: frozenset[str] = frozenset(), + ) -> None: self.rows: dict[str, _UserRow] = {} self.fail_ids = fail_ids self.commit_then_drop = commit_then_drop + self.raced_ids = raced_ids self.create_many_calls = 0 async def count(self, where: object = None) -> int: @@ -59,10 +67,15 @@ class _UserTable: rows = [_UserRow.model_validate(d) for d in data] if any(row.user_id in self.fail_ids for row in rows): raise RuntimeError("batch insert failed") + raced = [row.user_id for row in rows if row.user_id in self.raced_ids] + if raced: + for user_id in raced: + self.rows[user_id] = _UserRow(user_id=user_id, user_email=f"{user_id}@other-request.example") + raise UniqueViolationError({}, message="Unique constraint failed on the fields: (`user_id`)") for row in rows: self.rows[row.user_id] = row if self.commit_then_drop: - raise ConnectionError("connection reset after commit") + raise httpx.ReadError("connection reset after commit") return len(rows) async def update(self, where: dict[str, str], data: dict[str, object]) -> _UserRow: @@ -114,9 +127,13 @@ class _Tx: class _Db: def __init__( - self, teams: list[LiteLLM_TeamTable], fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False + self, + teams: list[LiteLLM_TeamTable], + fail_ids: frozenset[str] = frozenset(), + commit_then_drop: bool = False, + raced_ids: frozenset[str] = frozenset(), ) -> None: - self.litellm_usertable = _UserTable(fail_ids, commit_then_drop) + self.litellm_usertable = _UserTable(fail_ids, commit_then_drop, raced_ids) self.litellm_teamtable = _TeamTable(teams) self.litellm_teammembership = _MembershipTable() @@ -127,8 +144,9 @@ class _FakePrisma: teams: list[LiteLLM_TeamTable] | None = None, fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False, + raced_ids: frozenset[str] = frozenset(), ) -> None: - self.db = _Db(teams or [], fail_ids, commit_then_drop) + self.db = _Db(teams or [], fail_ids, commit_then_drop, raced_ids) self.tx_count = 0 self.locks: list[str] = [] @@ -276,6 +294,17 @@ async def test_insert_that_committed_but_lost_its_response_still_counts_as_creat assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1"] +@pytest.mark.asyncio +async def test_user_id_taken_by_a_concurrent_request_is_not_claimed_by_this_batch(): + prisma = _FakePrisma(teams=[_team("t1")], raced_ids=frozenset({"u1"})) + response = await _run(prisma, [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2", "teams": ["t1"]}]) + + assert [r.success for r in response.results] == [False, True] + assert "User id=u1 already exists" in (response.results[0].error or "") + assert prisma.db.litellm_usertable.rows["u1"].user_email == "u1@other-request.example" + assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u2"] + + @pytest.mark.asyncio async def test_team_write_failure_keeps_user_and_reports_it_on_the_row(): prisma = _FakePrisma(teams=[_team("t1"), _team("t2")])