diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 363f312336c..1f91eeedf64 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -21,7 +21,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_cache from litellm.proxy.utils import get_prisma_client_or_throw -from litellm.repositories.table_repositories import AccessGroupRepository +from litellm.repositories.table_repositories import AccessGroupRepository, TeamRepository from litellm.types.access_group import ( AccessGroupCreateRequest, AccessGroupResponse, @@ -75,11 +75,11 @@ class _AccessGroupTable(Protocol): class _TeamTable(Protocol): - async def find_unique(self, where: Mapping[str, object]) -> _TeamRecord | None: ... + async def find_unique(self, *, where: Mapping[str, object]) -> _TeamRecord | None: ... - async def find_many(self, where: Mapping[str, object]) -> Sequence[_TeamRecord]: ... + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_TeamRecord]: ... - async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... class _KeyTable(Protocol): @@ -153,17 +153,16 @@ async def _attached_team_ids_for( stored_team_ids: Final = tuple( dict.fromkeys(team_id for record in records for team_id in (record.assigned_team_ids or ())) ) - carrying: Final = {"access_group_ids": {"hasSome": group_ids}} # mutable-ok: prisma where must be a dict - listed: Final = {"team_id": {"in": stored_team_ids}} # mutable-ok: prisma where must be a dict - clauses: Final = (carrying, listed) if stored_team_ids else (carrying,) - where: Final = {"OR": clauses} # mutable-ok: prisma where must be a dict - return _attached_team_ids_by_group(records, await team_table.find_many(where=where)) + carrying: Final = {"access_group_ids": {"hasSome": group_ids}} # mutable-ok: prisma where is a dict + listed: Final = {"team_id": {"in": stored_team_ids}} # mutable-ok: prisma where is a dict + teams: Final = await team_table.find_many(where={"OR": (carrying, listed)}) # mutable-ok: prisma where is a dict + return _attached_team_ids_by_group(records, teams) async def _require_teams_exist(tx: _AccessGroupTx, team_ids: Sequence[str]) -> None: if not team_ids: return - where: Final = {"team_id": {"in": team_ids}} # mutable-ok: prisma where must be a dict + where: Final = {"team_id": {"in": team_ids}} # mutable-ok: prisma where is a dict found: Final = await tx.litellm_teamtable.find_many(where=where) missing: Final = frozenset(team_ids) - frozenset(team.team_id for team in found) if missing: @@ -441,7 +440,7 @@ async def list_access_groups( table: Final = AccessGroupRepository(prisma_client).table records: Final = await table.find_many(order={"created_at": "desc"}) - attached: Final = await _attached_team_ids_for(prisma_client.db.litellm_teamtable, records) + attached: Final = await _attached_team_ids_for(TeamRepository(prisma_client).table, records) return [_record_to_response(r, assigned_team_ids=attached[r.access_group_id]) for r in records] @@ -463,7 +462,7 @@ async def get_access_group( status_code=status.HTTP_404_NOT_FOUND, detail=f"Access group '{access_group_id}' not found", ) - attached: Final = await _attached_team_ids_for(prisma_client.db.litellm_teamtable, (record,)) + attached: Final = await _attached_team_ids_for(TeamRepository(prisma_client).table, (record,)) return _record_to_response(record, assigned_team_ids=attached[record.access_group_id]) @@ -516,7 +515,8 @@ async def update_access_group( ) await _require_teams_exist(tx, data.assigned_team_ids or ()) - old_team_ids: Final[set[str]] = set(existing.assigned_team_ids or []) + attached: Final = await _attached_team_ids_for(tx.litellm_teamtable, (existing,)) + old_team_ids: Final[set[str]] = set(attached[access_group_id]) old_key_ids: Final[set[str]] = set(existing.assigned_key_ids or []) new_team_ids: Final[set[str]] = ( set(update_fields["assigned_team_ids"] or []) if "assigned_team_ids" in update_fields else old_team_ids diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 18cf884f267..739d6ada71c 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -176,6 +176,10 @@ class PolicyAttachmentRepository(PrismaTableRepository["prisma_models.LiteLLM_Po table_name = "litellm_policyattachmenttable" +class TeamRepository(PrismaTableRepository["prisma_models.LiteLLM_TeamTable"]): + table_name = "litellm_teamtable" + + class DeletedTeamRepository(PrismaTableRepository["prisma_models.LiteLLM_DeletedTeamTable"]): table_name = "litellm_deletedteamtable" diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py index 39a6f78d14d..81816e21c10 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -317,8 +317,9 @@ def test_list_access_groups_attributes_teams_per_group_with_one_query(client_and assert body[1]["assigned_team_ids"] == ["team-y", "team-z"] mock_team_table.find_many.assert_awaited_once() - (carrying,) = mock_team_table.find_many.call_args.kwargs["where"]["OR"] + carrying, listed = mock_team_table.find_many.call_args.kwargs["where"]["OR"] assert list(carrying["access_group_ids"]["hasSome"]) == ["ag-1", "ag-2"] + assert list(listed["team_id"]["in"]) == [] @pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) @@ -1238,7 +1239,7 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks): mock_access_group_table.find_unique = AsyncMock(return_value=existing) team_to_remove = _make_team_record("team-remove", ["ag-update"]) - mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"])]) + mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"]), team_to_remove]) mock_team_table.find_unique = AsyncMock(return_value=team_to_remove) resp = client.put( @@ -1256,6 +1257,28 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks): assert "ag-update" not in call_kwargs["data"]["access_group_ids"] +def test_update_access_group_detaches_team_the_mirror_missed(client_and_mocks): + """Update removes the group from a team that carries it but was never written to the stored column.""" + client, mock_prisma, mock_access_group_table, *_ = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-keep"]) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + unmirrored = _make_team_record("team-unmirrored", ["ag-update", "ag-other"]) + mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"]), unmirrored]) + mock_team_table.find_unique = AsyncMock(return_value=unmirrored) + + resp = client.put("/v1/access_group/ag-update", json={"assigned_team_ids": ["team-keep"]}) + assert resp.status_code == 200 + + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-unmirrored"}) + mock_team_table.update.assert_awaited_once() + call_kwargs = mock_team_table.update.call_args.kwargs + assert call_kwargs["where"] == {"team_id": "team-unmirrored"} + assert call_kwargs["data"]["access_group_ids"] == ["ag-other"] + + def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_mocks): """Update does not sync teams when assigned_team_ids is absent from the payload.""" client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( @@ -1271,7 +1294,6 @@ def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_moc resp = client.put("/v1/access_group/ag-update", json={"description": "new desc"}) assert resp.status_code == 200 - mock_team_table.find_many.assert_not_awaited() mock_team_table.find_unique.assert_not_awaited() mock_team_table.update.assert_not_awaited() @@ -1405,7 +1427,7 @@ def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks) def test_update_access_group_null_assigned_ids_treated_as_empty(client_and_mocks): """Update with explicit null for assigned_*_ids clears the list and writes [] to DB.""" - client, mock_prisma, mock_table, *_ = client_and_mocks + client, _, mock_table, *_ = client_and_mocks existing = _make_access_group_record( access_group_id="ag-update", @@ -1425,4 +1447,3 @@ def test_update_access_group_null_assigned_ids_treated_as_empty(client_and_mocks update_call_kwargs = mock_table.update.call_args.kwargs assert update_call_kwargs["data"]["assigned_team_ids"] == [] assert update_call_kwargs["data"]["assigned_key_ids"] == [] - mock_prisma.db.litellm_teamtable.find_many.assert_not_awaited()