mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
fix(keys): write project assignment only while key row is still unassigned
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus
parent
1917700a57
commit
93d579fc9c
2 changed files with 101 additions and 6 deletions
|
|
@ -2467,6 +2467,47 @@ async def _apply_soft_budget_update(
|
|||
return remaining
|
||||
|
||||
|
||||
async def _write_guarded_project_assignment(
|
||||
table: "TableActions[prisma_models.LiteLLM_VerificationToken]",
|
||||
hashed_token: str,
|
||||
data: Mapping[str, object],
|
||||
) -> "prisma_models.LiteLLM_VerificationToken | None":
|
||||
updated_count: Final = await table.update_many(
|
||||
where={"token": hashed_token, "project_id": None},
|
||||
data=data,
|
||||
)
|
||||
if updated_count == 0:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Project reassignment is not supported. Use null to detach the key.",
|
||||
)
|
||||
return await table.find_unique(
|
||||
where={"token": hashed_token},
|
||||
include={"object_permission": True},
|
||||
)
|
||||
|
||||
|
||||
async def _update_key_row_assigning_project(
|
||||
prisma_client: PrismaClient,
|
||||
key: str,
|
||||
update_values: Mapping[str, object],
|
||||
) -> _KeyUpdateResult:
|
||||
"""Write a key update that assigns a project only while the row still has none."""
|
||||
hashed_token: Final = _hash_token_if_needed(key)
|
||||
updated_row: Final = await _write_guarded_project_assignment(
|
||||
table=_prisma_table(VerificationTokenRepository(prisma_client)),
|
||||
hashed_token=hashed_token,
|
||||
data=with_settings_updated_at(
|
||||
prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token}))
|
||||
),
|
||||
)
|
||||
updated_data: Final[Mapping[str, object]] = (
|
||||
updated_row.model_dump() if updated_row is not None else MappingProxyType({})
|
||||
)
|
||||
result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data}
|
||||
return result
|
||||
|
||||
|
||||
async def _update_key_row_with_soft_budget(
|
||||
prisma_client: PrismaClient,
|
||||
key: str,
|
||||
|
|
@ -2474,6 +2515,7 @@ async def _update_key_row_with_soft_budget(
|
|||
non_default_values: Mapping[str, object],
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
changed_by: str,
|
||||
expect_unassigned_project: bool,
|
||||
) -> _KeyUpdateResult:
|
||||
hashed_token: Final = _hash_token_if_needed(key)
|
||||
key_where: Final[_KeyRowWhere] = {"token": hashed_token}
|
||||
|
|
@ -2487,12 +2529,21 @@ async def _update_key_row_with_soft_budget(
|
|||
changed_by=changed_by,
|
||||
)
|
||||
include_object_permission: Final[prisma.types.LiteLLM_VerificationTokenInclude] = {"object_permission": True}
|
||||
updated_row: Final = await tx.litellm_verificationtoken.update(
|
||||
where=key_where,
|
||||
data=with_settings_updated_at(
|
||||
prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token}))
|
||||
),
|
||||
include=include_object_permission,
|
||||
update_payload: Final = with_settings_updated_at(
|
||||
prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token}))
|
||||
)
|
||||
updated_row: Final = (
|
||||
await _write_guarded_project_assignment(
|
||||
table=tx.litellm_verificationtoken,
|
||||
hashed_token=hashed_token,
|
||||
data=update_payload,
|
||||
)
|
||||
if expect_unassigned_project
|
||||
else await tx.litellm_verificationtoken.update(
|
||||
where=key_where,
|
||||
data=update_payload,
|
||||
include=include_object_permission,
|
||||
)
|
||||
)
|
||||
updated_data: Final[Mapping[str, object]] = (
|
||||
updated_row.model_dump() if updated_row is not None else MappingProxyType({})
|
||||
|
|
@ -3576,6 +3627,7 @@ async def update_key_fn(
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
is_project_assignment: Final = data.project_id is not None and existing_key_row.project_id is None
|
||||
response: Final = (
|
||||
await _update_key_row_with_soft_budget(
|
||||
prisma_client=prisma_client,
|
||||
|
|
@ -3584,8 +3636,15 @@ async def update_key_fn(
|
|||
non_default_values=update_values,
|
||||
existing_key_row=existing_key_row,
|
||||
changed_by=changed_by,
|
||||
expect_unassigned_project=is_project_assignment,
|
||||
)
|
||||
if "soft_budget" in data.model_fields_set
|
||||
else await _update_key_row_assigning_project(
|
||||
prisma_client=prisma_client,
|
||||
key=key,
|
||||
update_values=update_values,
|
||||
)
|
||||
if is_project_assignment
|
||||
else await prisma_client.update_data(token=key, data=MappingProxyType({**update_values, "token": key}))
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -61,6 +61,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
|
|||
_requested_end_user_budget_id,
|
||||
_save_deleted_verification_token_records,
|
||||
_transform_verification_tokens_to_deleted_records,
|
||||
_update_key_row_assigning_project,
|
||||
_validate_end_user_budget_id_change,
|
||||
_validate_max_budget,
|
||||
_validate_reset_spend_value,
|
||||
|
|
@ -20595,6 +20596,41 @@ async def test_prepare_key_update_data_includes_new_project_assignment():
|
|||
assert result["project_id"] == "project-orbit"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_assignment_write_requires_row_still_unassigned():
|
||||
database = MagicMock()
|
||||
database.jsonify_object = lambda data: dict(data)
|
||||
database.db.litellm_verificationtoken.update_many = AsyncMock(return_value=1)
|
||||
row = MagicMock()
|
||||
row.model_dump = MagicMock(return_value={"project_id": "project-orbit"})
|
||||
database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=row)
|
||||
|
||||
result: Final = await _update_key_row_assigning_project(
|
||||
prisma_client=database, key="sk-assign", update_values={"project_id": "project-orbit"}
|
||||
)
|
||||
|
||||
update_where: Final = database.db.litellm_verificationtoken.update_many.await_args.kwargs["where"]
|
||||
assert update_where["project_id"] is None
|
||||
assert update_where["token"] == result["token"]
|
||||
assert result["data"] == {"project_id": "project-orbit"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_assignment_write_rejects_when_row_gained_project():
|
||||
database = MagicMock()
|
||||
database.jsonify_object = lambda data: dict(data)
|
||||
database.db.litellm_verificationtoken.update_many = AsyncMock(return_value=0)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _update_key_row_assigning_project(
|
||||
prisma_client=database, key="sk-assign", update_values={"project_id": "project-orbit"}
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 400
|
||||
assert "reassignment" in str(exc.value.detail)
|
||||
database.db.litellm_verificationtoken.find_unique.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_creator_cannot_detach_project_without_admin_access():
|
||||
existing: Final = LiteLLM_VerificationToken(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue