fix(users): preserve retries for existing-user upserts

litellm/proxy/utils.py

@@ -310,6 +310,7 @@ if TYPE_CHECKING:
     from prisma import models as prisma_models
     from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions
     from prisma.client import TransactionManager
+    from prisma.errors import UniqueViolationError
     from prisma.models import LiteLLM_DeprecatedVerificationToken
     from prisma.types import HttpConfig, LiteLLM_VerificationTokenInclude
 
@@ -5536,6 +5537,15 @@ def _user_unique_conflict_message(metadata: object, user_id: object) -> str:
             return "User conflicts with an existing unique value"
 
 
+def _raise_user_unique_conflict(error: "UniqueViolationError", user_id: object, *, allow_existing: bool) -> NoReturn:
+    if allow_existing:
+        raise error
+    raise HTTPException(
+        status_code=409,
+        detail={"error": _user_unique_conflict_message(error.meta, user_id)},
+    ) from error
+
+
 class PrismaClient:
     spend_log_transactions: list = []
     _spend_log_transactions_lock = asyncio.Lock()
@@ -6480,10 +6490,7 @@ class PrismaClient:
                             detail="SSO for more than 5 users requires a LiteLLM Enterprise license. Existing users can still sign in.",
                         )
                 except UniqueViolationError as e:
-                    raise HTTPException(
-                        status_code=409,
-                        detail={"error": _user_unique_conflict_message(e.meta, db_data["user_id"])},
-                    ) from e
+                    _raise_user_unique_conflict(e, db_data["user_id"], allow_existing=allow_existing_user)
                 except Exception as e:
                     if (
                         "Foreign key constraint failed on the field: `LiteLLM_UserTable_organization_id_fkey (index)`"

tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py

@@ -221,6 +221,20 @@ async def test_user_conflict_names_constraint_and_does_not_retry(
     assert create.await_count == 1
 
 
+@pytest.mark.asyncio
+async def test_user_upsert_recovers_from_concurrent_creation(
+    prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    row: Final = LiteLLM_UserTable(user_id="requested-user")
+    conflict: Final = UniqueViolationError({"user_facing_error": {"meta": {"target": ["user_id"]}}})
+    prisma_client.db.litellm_usertable.upsert = AsyncMock(side_effect=[conflict, row])
+    monkeypatch.setattr("asyncio.sleep", AsyncMock())
+
+    result: Final = await prisma_client.insert_data(data={"user_id": row.user_id}, table_name="user")
+
+    assert result is row
+
+
 @pytest.mark.asyncio
 async def test_user_insert_still_retries_transient_connection_failure(
     prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch