scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-02 07:13:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 148

fix(proxy): look up organization members by email with find_many

Browse source 
user_email is not @unique in LiteLLM_UserTable, so find_unique on it always
raises in Prisma. /organization/member_add returned 500 for every request that
resolved a member by email, and member_update/member_delete rejected every
email with 400.

Use find_many in add_member_to_organization and find_member_if_email, and let
the duplicate-email HTTPException reach the endpoint as a 400 instead of being
rewrapped into a 500.

Fixes #44766
This commit is contained in:
Alex Castillo 2026-10-06 00:22:21 -04:00
parent b8239a9873
commit 45cd42ed1a
2 changed files with 158 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -119,6 +119,8 @@ class _ObjectPermissionRow(Protocol):
class _UserTableClient(Protocol):
async def find_unique(self, where: Mapping[str, object]) -> "PrismaUserTable | None": ...
async def find_many(self, where: Mapping[str, object]) -> "Sequence[PrismaUserTable]": ...
class _BudgetTableClient(Protocol):
async def create(self, data: Mapping[str, object]) -> "PrismaBudgetTable": ...
@ -1414,16 +1416,12 @@ async def find_member_if_email(user_email: str, prisma_client: PrismaClient) ->
"error": f"Unique user not found for user_email={user_email}. Potential duplicate OR non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead."
},
)
try:
existing_user_email_row: Final = await UserRepository(prisma_client).table.find_unique(
where={"user_email": user_email}
)
except Exception:
existing_user_email_rows: Final = await _table(UserRepository(prisma_client)).find_many(
where={"user_email": user_email}
)
if len(existing_user_email_rows) != 1:
raise not_unique_user_email_error
if existing_user_email_row is None:
raise not_unique_user_email_error
existing_user_email_row_pydantic: Final = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump())
return existing_user_email_row_pydantic
return LiteLLM_UserTable.model_validate(existing_user_email_rows[0].model_dump())
@router.patch(
@ -1643,25 +1641,20 @@ async def add_member_to_organization(
try:
user_object: LiteLLM_UserTable | None = None
existing_user_id_row = None
existing_user_email_row = None
## Check if user exists in LiteLLM_UserTable - user exists - either the user_id or user_email is in LiteLLM_UserTable
if member.user_id is not None:
existing_user_id_row = await _table(UserRepository(prisma_client)).find_unique(
where={"user_id": member.user_id}
)
if existing_user_id_row is None and member.user_email is not None:
try:
existing_user_email_row = await UserRepository(prisma_client).table.find_unique(
where={"user_email": member.user_email}
)
except Exception as e:
raise ValueError(
f"Potential NON-Existent or Duplicate user email in DB: Error finding a unique instance of user_email={member.user_email} in LiteLLM_UserTable.: {e}"
)
existing_user_email_rows: Final = (
await _table(UserRepository(prisma_client)).find_many(where={"user_email": member.user_email})
if existing_user_id_row is None and member.user_email is not None
else ()
)
## If user does not exist, create a new user
if existing_user_id_row is None and existing_user_email_row is None:
if existing_user_id_row is None and not existing_user_email_rows:
# Create a new user - since user does not exist
user_id: Final[str] = member.user_id or str(uuid.uuid4())
new_user_defaults: Final = get_new_internal_user_defaults(
@ -1672,16 +1665,13 @@ async def add_member_to_organization(
_returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user")
if _returned_user is not None:
user_object = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
elif existing_user_email_row is not None and (
len(existing_user_email_row) # pyright: ignore[reportArgumentType] # find_unique yields a row, not a list
> 1
):
elif len(existing_user_email_rows) > 1:
raise HTTPException(
status_code=400,
detail={"error": "Multiple users with this email found in db. Please use 'user_id' instead."},
)
elif existing_user_email_row is not None:
user_object = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump())
elif existing_user_email_rows:
user_object = LiteLLM_UserTable.model_validate(existing_user_email_rows[0].model_dump())
elif existing_user_id_row is not None:
user_object = LiteLLM_UserTable.model_validate(existing_user_id_row.model_dump())
else:
@ -1708,5 +1698,7 @@ async def add_member_to_organization(
)
return user_object, organization_membership
except HTTPException:
raise
except Exception as e:
raise ValueError(f"Error adding member={member} to organization={organization_id}: {e}")

140
tests/unit/proxy/management_endpoints/test_organization_endpoints.py
View file

@ -1573,3 +1573,143 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke
assert all(cache.get_cache(key=cache_key) is None for cache_key in doomed_cache_keys)
assert all(cache.get_cache(key=cache_key) == {"retained": True} for cache_key in kept_cache_keys)
assert jwt_table.rows == (kept_row,)
class _PrismaLikeUserTable:
"""Rejects find_unique on user_email like Prisma does, because the column is not @unique"""
def __init__(self, rows: tuple) -> None:
self.rows = rows
async def find_unique(self, where: Mapping[str, object]):
if set(where) != {"user_id"}:
raise ValueError(f"Could not find field at `findUniqueLiteLLM_UserTable.where.{next(iter(where))}`")
return next((row for row in self.rows if row.user_id == where["user_id"]), None)
async def find_many(self, where: Mapping[str, object]):
return [row for row in self.rows if all(getattr(row, k) == v for k, v in where.items())]
def _org_prisma_with_users(monkeypatch, users: tuple):
from datetime import datetime
from litellm.proxy._types import LiteLLM_OrganizationMembershipTable, LiteLLM_UserTable
async def create_membership(data):
return LiteLLM_OrganizationMembershipTable(
user_id=data["user_id"],
organization_id=data["organization_id"],
user_role=data["user_role"],
created_at=datetime(2024, 1, 1),
updated_at=datetime(2024, 1, 1),
)
async def insert_user(data, table_name):
return LiteLLM_UserTable(user_id=data["user_id"], user_email=data.get("user_email"), user_role="internal_user")
mock_prisma = SimpleNamespace(
db=SimpleNamespace(
litellm_organizationtable=SimpleNamespace(find_unique=AsyncMock(return_value=SimpleNamespace())),
litellm_usertable=_PrismaLikeUserTable(users),
litellm_organizationmembership=SimpleNamespace(create=create_membership),
),
insert_data=insert_user,
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.organization_endpoints._verify_org_access",
AsyncMock(),
)
return mock_prisma
async def _add_org_member(member: dict):
from litellm.proxy._types import LitellmUserRoles, OrganizationMemberAddRequest, UserAPIKeyAuth
from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add
return await organization_member_add(
data=OrganizationMemberAddRequest(organization_id="org-1", member=member),
http_request=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"member, expected_user_id, expected_email",
[
({"user_email": "alice@example.com", "role": "internal_user"}, "alice-id", "alice@example.com"),
(
{"user_id": "new-id", "user_email": "alice@example.com", "role": "internal_user"},
"alice-id",
"alice@example.com",
),
({"user_id": "new-id", "user_email": "bob@example.com", "role": "org_admin"}, "new-id", "bob@example.com"),
({"user_email": "bob@example.com", "role": "internal_user"}, None, "bob@example.com"),
],
ids=["email-of-existing-user", "unknown-id-with-existing-email", "unknown-id-with-new-email", "new-email-only"],
)
async def test_organization_member_add_resolves_member_by_email(member, expected_user_id, expected_email, monkeypatch):
from litellm.proxy._types import LiteLLM_UserTable
alice: Final = LiteLLM_UserTable(user_id="alice-id", user_email="alice@example.com", user_role="internal_user")
_org_prisma_with_users(monkeypatch, (alice,))
response = await _add_org_member(member)
user = response.updated_users[0]
membership = response.updated_organization_memberships[0]
assert user.user_email == expected_email, response
assert membership.user_id == user.user_id, response
assert membership.user_role == member["role"], response
if expected_user_id is None:
assert user.user_id != alice.user_id, response
else:
assert user.user_id == expected_user_id, response
@pytest.mark.asyncio
async def test_organization_member_add_duplicate_email_returns_400(monkeypatch):
from litellm.proxy._types import LiteLLM_UserTable, ProxyException
_org_prisma_with_users(
monkeypatch,
tuple(
LiteLLM_UserTable(user_id=user_id, user_email="dup@example.com", user_role="internal_user")
for user_id in ("dup-1", "dup-2")
),
)
with pytest.raises(ProxyException) as exc:
await _add_org_member({"user_email": "dup@example.com", "role": "internal_user"})
assert int(exc.value.code) == 400, exc.value.message
assert "Multiple users with this email" in str(exc.value.message)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"rows, expected_user_id",
[
((("alice-id", "alice@example.com"),), "alice-id"),
((), None),
((("a", "alice@example.com"), ("b", "alice@example.com")), None),
],
ids=["unique", "missing", "duplicate"],
)
async def test_find_member_if_email_resolves_only_a_unique_email(rows, expected_user_id, monkeypatch):
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.management_endpoints.organization_endpoints import find_member_if_email
prisma = _org_prisma_with_users(
monkeypatch,
tuple(
LiteLLM_UserTable(user_id=user_id, user_email=email, user_role="internal_user") for user_id, email in rows
),
)
if expected_user_id is None:
with pytest.raises(HTTPException) as exc:
await find_member_if_email("alice@example.com", prisma)
assert exc.value.status_code == 400
else:
assert (await find_member_if_email("alice@example.com", prisma)).user_id == expected_user_id