diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c8895c01f7b..d28401524ed 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -5482,7 +5482,9 @@ async def generate_key_helper_fn( if query_type == "insert_data": user_row = cast( # cast-ok: table_name="user" is the insert_data branch returning the user row "prisma_models.LiteLLM_UserTable | None", - await prisma_client.insert_data(data=user_data, table_name="user"), + await prisma_client.insert_data( + data=user_data, table_name="user", allow_existing_user=request_type != "user" + ), ) if user_row is None: diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index e48f68a2d0a..66c2e54f5c6 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1294,7 +1294,7 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou if litellm.default_internal_user_params: default_role = litellm.default_internal_user_params.get("user_role") - new_user_request: Final = NewUserRequest( + new_user_request: Final = SCIMNewUserRequest( user_id=user_id, user_email=user_id, # We don't have email from group membership user_alias=None, @@ -1873,7 +1873,7 @@ async def create_user( admin_group: Final = await _get_scim_admin_group() resolved_role: Final = _resolve_scim_user_role(user.groups or [], admin_group, default_role) - new_user_request: Final = NewUserRequest( + new_user_request: Final = SCIMNewUserRequest( user_id=user_id, user_email=user_data["user_email"], user_alias=user_data["user_alias"], diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 53079d84fdd..aaa261eb19b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -78,7 +78,6 @@ from litellm.proxy._types import ( LitellmUserRoles, Member, NewTeamRequest, - NewUserRequest, NewUserRequestTeam, NewUserResponse, ProxyErrorTypes, @@ -164,6 +163,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( MicrosoftGraphAPIUserGroupResponse, MicrosoftServicePrincipalTeam, RoleMappings, + SSONewUserRequest, TeamMappings, ) from litellm.types.proxy.ui_sso import ParsedOpenIDResult @@ -2915,7 +2915,7 @@ async def insert_sso_user( if user_defined_values["user_role"] is None: user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY - new_user_request: Final = NewUserRequest( + new_user_request: Final = SSONewUserRequest( user_id=user_defined_values["user_id"], user_email=normalize_email(user_defined_values["user_email"]), user_alias=_get_sso_user_alias(result_openid), @@ -2924,11 +2924,9 @@ async def insert_sso_user( budget_duration=user_defined_values["budget_duration"], sso_user_id=user_defined_values["user_id"], auto_create_key=False, + metadata={"auth_provider": getattr(result_openid, "provider")} if hasattr(result_openid, "provider") else {}, ) - if result_openid and hasattr(result_openid, "provider"): - new_user_request.metadata = {"auth_provider": getattr(result_openid, "provider")} - from litellm.proxy.proxy_server import premium_user with free_sso_user_creation(new_user_request.user_id, enabled=premium_user is not True): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 00187790901..3102674376a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -313,6 +313,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 @@ -5559,6 +5560,25 @@ _LOCK_TEAMS_BY_ID_SQL: Final = ( ) +def _user_unique_conflict_message(metadata: object, user_id: object) -> str: + match metadata: + case {"target": ["user_id"]}: + return f"User with id {user_id} already exists" + case {"target": ["sso_user_id"]}: + return "User with this sso_user_id already exists" + case _: + 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() @@ -6446,12 +6466,15 @@ class PrismaClient: max_tries=3, # maximum number of retries max_time=10, # maximum total time to retry for on_backoff=on_backoff, # specifying the function to call on backoff + giveup=lambda exc: isinstance(exc, HTTPException) and exc.status_code == 409, ) @log_db_metrics async def insert_data( self, data: Mapping[str, object], table_name: Literal["user", "key", "config", "spend", "team", "user_notification"], + *, + allow_existing_user: bool = True, ): """ Add a key to the database. If it already exists, do nothing. @@ -6487,14 +6510,20 @@ class PrismaClient: verbose_proxy_logger.info("Data Inserted into Keys Table") return new_verification_token elif table_name == "user": + from prisma.errors import UniqueViolationError + db_data = self.jsonify_object(data=data) try: - new_user_row: Final = await UserRepository(self).upsert_user(db_data) + new_user_row: Final = await UserRepository(self).upsert_user( + db_data, allow_existing=allow_existing_user + ) if new_user_row is None: raise HTTPException( status_code=403, detail="SSO for more than 5 users requires a LiteLLM Enterprise license. Existing users can still sign in.", ) + except UniqueViolationError as 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)`" diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 1fce0717ea3..28fba2de564 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -168,7 +168,9 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): """ return await _count_billable_users(self.table) - async def upsert_user(self, data: Mapping[str, object]) -> "prisma_models.LiteLLM_UserTable | None": + async def upsert_user( + self, data: Mapping[str, object], *, allow_existing: bool = True + ) -> "prisma_models.LiteLLM_UserTable | None": """Return the user row, or None when a new free SSO user would exceed capacity.""" scope: Final = _FREE_SSO_CREATION.get() if ( @@ -176,6 +178,8 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): or (scope[0] is not None and scope[0] != data["user_id"]) or scope[1] is not asyncio.current_task() ): + if not allow_existing: + return await self.table.create(data=dict(data)) return await self.table.upsert( where={"user_id": data["user_id"]}, data={"create": dict(data), "update": {}} ) diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index fc0b73f77b6..e79b509da1b 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -120,6 +120,7 @@ class BulkNewUserItem(NewUserRequest): model_config = ConfigDict(extra="forbid", protected_namespaces=()) + user_id: str | None = None auto_create_key: bool = False @field_validator("send_invite_email") diff --git a/litellm/types/proxy/management_endpoints/scim_v2.py b/litellm/types/proxy/management_endpoints/scim_v2.py index fadb309e0cd..72a39c61967 100644 --- a/litellm/types/proxy/management_endpoints/scim_v2.py +++ b/litellm/types/proxy/management_endpoints/scim_v2.py @@ -13,6 +13,7 @@ from pydantic import ( from pydantic_core.core_schema import SerializerFunctionWrapHandler from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.proxy.management_endpoints.user_requests import NewUserRequest SCIM_ENTERPRISE_USER_SCHEMA: Final = "urn:ietf:params:scim:schemas:extension:enterprise:2.0:User" SCIM_ENTERPRISE_METADATA_KEY: Final = "scim_enterprise" @@ -26,6 +27,14 @@ _SCIM_SERIALIZED_FIELDS_ADAPTER: Final = TypeAdapter( ) +class SCIMNewUserRequest(NewUserRequest): + """Preserve opaque directory IDs when provisioning SCIM users and group members.""" + + model_config = ConfigDict(frozen=True) + + user_id: str | None = None + + class LiteLLM_UserScimMetadata(LiteLLMBaseModel): """ Scim metadata stored in LiteLLM_UserTable.metadata diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index affb6d2e9ce..38ed799f38b 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -1,12 +1,19 @@ from typing import Literal -from pydantic import Field +from pydantic import ConfigDict, Field from typing_extensions import TypedDict from litellm.types.proxy.auth.user_roles import KeyManagementRoutes, LitellmUserRoles +from litellm.types.proxy.management_endpoints.user_requests import NewUserRequest from litellm.types.utils import LiteLLMPydanticObjectBase +class SSONewUserRequest(NewUserRequest): + model_config = ConfigDict(frozen=True) + + user_id: str | None = None + + class LiteLLM_UpperboundKeyGenerateParams(LiteLLMPydanticObjectBase): """ Set default upperbound to max budget a key called via `/key/generate` can be. diff --git a/litellm/types/proxy/management_endpoints/user_requests.py b/litellm/types/proxy/management_endpoints/user_requests.py index e53dc4152c3..edda1f78430 100644 --- a/litellm/types/proxy/management_endpoints/user_requests.py +++ b/litellm/types/proxy/management_endpoints/user_requests.py @@ -1,7 +1,7 @@ from collections.abc import Mapping -from typing import Literal +from typing import Annotated, Final, Literal -from pydantic import Field, field_validator, model_validator +from pydantic import AfterValidator, Field, field_validator, model_validator from litellm.models.user import LiteLLM_UserTable from litellm.types.proxy.auth.user_roles import LitellmUserRoles @@ -15,7 +15,17 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase): user_role: Literal["user", "admin"] = "user" +def _validate_custom_user_id(value: str) -> str: + normalized: Final = value.strip() + if not normalized: + raise ValueError("user_id must not be empty or contain only whitespace") + if "\x00" in normalized: + raise ValueError("user_id must not contain NUL characters") + return normalized + + class NewUserRequest(GenerateRequestBase): + user_id: Annotated[str, AfterValidator(_validate_custom_user_id)] | None = None tpm_limit: int | None = Field(default=None, ge=0) rpm_limit: int | None = Field(default=None, ge=0) max_budget: float | None = None diff --git a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 180362f45f1..94e0f91ecd1 100644 --- a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -411,10 +411,11 @@ async def test_create_user_without_groups_defers_to_default_team(mocker: MockerF @pytest.mark.asyncio -async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch): +@pytest.mark.parametrize("user_id", ["new-user", " employee-42 "]) +async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch, user_id: str): scim_user = SCIMUser( schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], - userName="new-user", + userName=user_id, emails=[SCIMUserEmail(value="new@example.com")], groups=[SCIMUserGroup(value="idp-team")], ) @@ -428,10 +429,12 @@ async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, mo await create_user(user=scim_user) assert new_user_mock.call_args.kwargs["data"].teams == ["idp-team"] + assert new_user_mock.call_args.kwargs["data"].user_id == user_id @pytest.mark.asyncio -async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch): +@pytest.mark.parametrize("user_id", ["group-user", " employee-42 "]) +async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch, user_id: str): monkeypatch.setattr( "litellm.default_internal_user_params", {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, @@ -439,14 +442,15 @@ async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFi ) new_user_mock = mocker.patch( # test-quality-ok: new_user is imported inside the helper, not injectable "litellm.proxy.management_endpoints.internal_user_endpoints.new_user", - AsyncMock(return_value=NewUserResponse(user_id="group-user", key="k")), + AsyncMock(return_value=NewUserResponse(user_id=user_id, key="k")), ) - created = await _create_user_if_not_exists(user_id="group-user") + created = await _create_user_if_not_exists(user_id=user_id) assert created is not None assert new_user_mock.call_args.kwargs["data"].teams is None assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + assert new_user_mock.call_args.kwargs["data"].user_id == user_id @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index cee0c90f80c..b96a5b606eb 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -13,6 +13,7 @@ import httpx import pytest from fastapi import HTTPException from fastapi.testclient import TestClient +from pydantic import ValidationError from pytest_mock import MockerFixture from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -48,6 +49,51 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import ( client = TestClient(app) + +@pytest.mark.parametrize("user_id", ["", " \t\n", "user\x00id", 123, True, [], {}]) +def test_new_user_rejects_invalid_custom_id(user_id: object) -> None: + with pytest.raises(ValidationError) as exc: + NewUserRequest.model_validate({"user_id": user_id}) + assert exc.value.errors()[0]["loc"] == ("user_id",) + + +@pytest.mark.parametrize("user_id", [None, "employee-42", "社員@example.test", " employee-42 "]) +def test_new_user_normalizes_custom_id_and_preserves_auto_generation(user_id: str | None) -> None: + request: Final = NewUserRequest(user_id=user_id) + assert request.user_id == (user_id.strip() if user_id is not None else None) + + +@pytest.mark.asyncio +async def test_new_user_returns_conflict_when_id_is_taken_after_duplicate_check(mocker: MockerFixture) -> None: + from prisma.errors import UniqueViolationError + + from litellm.proxy import proxy_server + from litellm.proxy.utils import PrismaClient + + table: Final = SimpleNamespace( + find_first=AsyncMock(return_value=None), + count=AsyncMock(return_value=0), + create=AsyncMock(side_effect=UniqueViolationError( + {"user_facing_error": {"meta": {"target": ["user_id"]}}}, + message="Unique constraint failed on user_id", + )), + upsert=AsyncMock(return_value=LiteLLM_UserTable(user_id="employee-42")), + ) + mocker.patch("prisma.Prisma", return_value=SimpleNamespace(litellm_usertable=table)) + database: Final = PrismaClient( + database_url="postgresql://unused:unused@localhost/unused", + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", database) + with pytest.raises(ProxyException) as exc: + await new_user( + data=NewUserRequest(user_id="employee-42", auto_create_key=False, send_invite_email=False), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert str(exc.value.code) == "409" + assert "employee-42" in exc.value.message + + if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache from litellm.proxy.utils import PrismaClient @@ -4862,7 +4908,9 @@ async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget, models = [] class _FakePrisma: - async def insert_data(self, data, table_name): + async def insert_data( + self, data: Mapping[str, object], table_name: str, *, allow_existing_user: bool = True + ) -> _FakeUserRow | SimpleNamespace: if table_name == "user": captured["user_data"] = dict(data) return _FakeUserRow() diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index ba78fefd653..6ccb019b937 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ( LiteLLM_JWTAuth, LiteLLM_UserTable, NewUserResponse, + SSOUserDefinedValues, ProxyException, ) from litellm.proxy.auth.handle_jwt import JWTHandler @@ -52,6 +53,62 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( from tests._master_key import MASTER_KEY +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("billable_users", "premium", "admitted"), + [(4, False, True), (5, False, False), (26, False, False), (26, True, True)], +) +async def test_sso_auto_creation_enforces_free_user_cap_through_new_user( + monkeypatch: pytest.MonkeyPatch, billable_users: int, premium: bool, admitted: bool +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.utils import PrismaClient + + row: Final = LiteLLM_UserTable(user_id="new-sso-user", user_email="new-sso-user@example.test") + table: Final = SimpleNamespace( + find_first=AsyncMock(return_value=None), + find_unique=AsyncMock(return_value=None), + count=AsyncMock(side_effect=lambda where=None: billable_users if where is None else 0), + create=AsyncMock(return_value=row), + upsert=AsyncMock(return_value=row), + ) + transaction: Final = MagicMock() + transaction.__aenter__ = AsyncMock( + return_value=SimpleNamespace(litellm_usertable=table, execute_raw=AsyncMock(), query_raw=AsyncMock()) + ) + transaction.__aexit__ = AsyncMock(return_value=None) + monkeypatch.setattr( + "prisma.Prisma", + lambda **kwargs: SimpleNamespace(litellm_usertable=table, tx=MagicMock(return_value=transaction)), + ) + database: Final = PrismaClient( + database_url="postgresql://unused:unused@localhost/unused", + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "premium_user", premium) + monkeypatch.setattr("asyncio.sleep", AsyncMock()) + values: Final[SSOUserDefinedValues] = { + "models": [], + "user_id": row.user_id, + "user_email": row.user_email, + "user_role": "internal_user_viewer", + "max_budget": None, + "budget_duration": None, + } + if admitted: + response: Final = await insert_sso_user(CustomOpenID(id=row.user_id, team_ids=[]), values) + assert response.user_id == row.user_id + assert table.create.await_count + table.upsert.await_count == 1 + else: + with pytest.raises(ProxyException) as error: + await insert_sso_user(CustomOpenID(id=row.user_id, team_ids=[]), values) + assert str(error.value.code) == "403" + assert "SSO for more than 5 users" in error.value.message + table.create.assert_not_awaited() + table.upsert.assert_not_awaited() + + def _sso_http_handler(route: Callable[[httpx.Request], httpx.Response] | None = None) -> AsyncHTTPHandler: """litellm's SSO httpx client answered in-process, by ``route`` or by the active respx mock.""" handler = AsyncHTTPHandler() @@ -1281,7 +1338,8 @@ async def test_upsert_sso_user_fills_user_alias_for_existing_user(): @pytest.mark.asyncio -async def test_insert_sso_user_sets_user_alias_from_display_name(): +@pytest.mark.parametrize("user_id", ["S-1-5-21-adfs-user", " employee-42 "]) +async def test_insert_sso_user_sets_user_alias_from_display_name(user_id: str): """ A newly created SSO user is inserted with the IdP display name as user_alias. """ @@ -1290,7 +1348,7 @@ async def test_insert_sso_user_sets_user_alias_from_display_name(): from litellm.proxy.management_endpoints.ui_sso import insert_sso_user sso_result = CustomOpenID( - id="S-1-5-21-adfs-user", + id=user_id, email="jane.doe@example.com", first_name="Jane", last_name="Doe", @@ -1300,7 +1358,7 @@ async def test_insert_sso_user_sets_user_alias_from_display_name(): ) user_defined_values: SSOUserDefinedValues = { "models": [], - "user_id": "S-1-5-21-adfs-user", + "user_id": user_id, "user_email": "jane.doe@example.com", "max_budget": None, "user_role": "internal_user", @@ -1309,12 +1367,14 @@ async def test_insert_sso_user_sets_user_alias_from_display_name(): with patch( "litellm.proxy.management_endpoints.ui_sso.new_user", - return_value=NewUserResponse(user_id="S-1-5-21-adfs-user", key="sk-xxxxx", teams=None), + return_value=NewUserResponse(user_id=user_id, key="sk-xxxxx", teams=None), ) as mock_new_user: await insert_sso_user(result_openid=sso_result, user_defined_values=user_defined_values) new_user_request = mock_new_user.call_args.kwargs["data"] - assert new_user_request.user_id == "S-1-5-21-adfs-user" + assert new_user_request.user_id == user_id + assert new_user_request.sso_user_id == user_id + assert new_user_request.metadata == {"auth_provider": "generic"} assert new_user_request.user_email == "jane.doe@example.com" assert new_user_request.user_alias == "Doe, Jane" diff --git a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py index cf1ac61c93a..2b966a8ebeb 100644 --- a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py @@ -191,6 +191,19 @@ async def _run(prisma, users, caller=ADMIN, license=None, generate_key=_no_keys) ) +@pytest.mark.asyncio +@pytest.mark.parametrize("user_id", [None, "", " employee-42 ", " \t"]) +async def test_bulk_keeps_existing_user_id_generation_and_whitespace_behavior(user_id: str | None) -> None: + prisma: Final = _FakePrisma() + response: Final = await _run(prisma, [{"user_id": user_id}]) + result: Final = response.data[0] + assert result.success + assert result.user_id + assert result.user_id in prisma.db.litellm_usertable.rows + if user_id: + assert result.user_id == user_id + + @pytest.mark.asyncio async def test_creates_users_and_team_membership_in_every_store(): prisma = _FakePrisma(teams=[_team("t1", [Member(user_id="existing", role="admin")]), _team("t2")]) diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py index e335b08f040..8603de0c3a1 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_writes.py @@ -18,8 +18,10 @@ from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException +from prisma.errors import UniqueViolationError from litellm._service_logger import ServiceTypes from litellm.models.user import LiteLLM_UserTable @@ -192,6 +194,62 @@ async def test_insert_data_team_serializes_members(prisma_client: PrismaClient) assert create_payload["team_id"] == "t1" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata", "message"), + [ + ({"target": ["user_id"]}, "User with id requested-user already exists"), + ({"target": ["sso_user_id"]}, "User with this sso_user_id already exists"), + (None, "User conflicts with an existing unique value"), + ({"target": ["unknown_field"]}, "User conflicts with an existing unique value"), + ], +) +async def test_user_conflict_names_constraint_and_does_not_retry( + prisma_client: PrismaClient, + metadata: Mapping[str, object] | None, + message: str, +) -> None: + create: Final = AsyncMock(side_effect=UniqueViolationError({"user_facing_error": {"meta": metadata}})) + prisma_client.db.litellm_usertable.create = create + with pytest.raises(HTTPException) as error: + await prisma_client.insert_data( + data={"user_id": "requested-user", "sso_user_id": "sso-subject"}, + table_name="user", allow_existing_user=False, + ) + assert error.value.status_code == 409 + assert error.value.detail == {"error": message} + 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 +) -> None: + row: Final = LiteLLM_UserTable(user_id="requested-user") + create: Final = AsyncMock(side_effect=[httpx.ConnectError("connection unavailable"), row]) + prisma_client.db.litellm_usertable.create = create + monkeypatch.setattr("asyncio.sleep", AsyncMock()) + result: Final = await prisma_client.insert_data( + data={"user_id": row.user_id}, table_name="user", allow_existing_user=False + ) + assert result is row + assert create.await_count == 2 + + @pytest.mark.asyncio async def test_insert_data_user_organization_fk_raises_400( prisma_client: PrismaClient,