feat(BA-4877): add service/repository layer for bulk role assignment and revocation

changes/9722.feature.md

@@ -0,0 +1 @@
+Add service/repository layer for bulk role assignment and revocation

src/ai/backend/manager/data/permission/role.py

@@ -148,6 +148,55 @@ class UserRoleRevocationData:
     role_id: uuid.UUID
 
 
+@dataclass(frozen=True)
+class BulkUserRoleAssignmentInput:
+    """Input for bulk assigning a role to multiple users."""
+
+    role_id: uuid.UUID
+    user_ids: list[uuid.UUID]
+    granted_by: uuid.UUID | None = None
+
+
+@dataclass(frozen=True)
+class BulkRoleAssignmentFailure:
+    """Failure information for a single user in bulk role assignment."""
+
+    user_id: uuid.UUID
+    message: str
+
+
+@dataclass(frozen=True)
+class BulkRoleAssignmentResultData:
+    """Result of bulk role assignment."""
+
+    successes: list[UserRoleAssignmentData] = field(default_factory=list)
+    failures: list[BulkRoleAssignmentFailure] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class BulkUserRoleRevocationInput:
+    """Input for bulk revoking a role from multiple users."""
+
+    role_id: uuid.UUID
+    user_ids: list[uuid.UUID]
+
+
+@dataclass(frozen=True)
+class BulkRoleRevocationFailure:
+    """Failure information for a single user in bulk role revocation."""
+
+    user_id: uuid.UUID
+    message: str
+
+
+@dataclass(frozen=True)
+class BulkRoleRevocationResultData:
+    """Result of bulk role revocation."""
+
+    successes: list[UserRoleRevocationData] = field(default_factory=list)
+    failures: list[BulkRoleRevocationFailure] = field(default_factory=list)
+
+
 @dataclass(frozen=True)
 class RoleListResult(SearchResult[RoleData]):
     """Result of role search with pagination info."""

src/ai/backend/manager/repositories/permission_controller/db_source/db_source.py

@@ -1,3 +1,4 @@
+import logging
 import uuid
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
@@ -10,6 +11,7 @@
 from ai.backend.common.data.permission.types import (
     RelationType,
 )
+from ai.backend.logging.utils import BraceStyleAdapter
 from ai.backend.manager.data.permission.entity import (
     ElementAssociationListResult,
     EntityData,
@@ -26,9 +28,13 @@
 from ai.backend.manager.data.permission.role import (
     AssignedUserData,
     AssignedUserListResult,
+    BulkRoleRevocationFailure,
+    BulkRoleRevocationResultData,
+    BulkUserRoleRevocationInput,
     RoleListResult,
     RolePermissionsUpdateInput,
     UserRoleAssignmentInput,
+    UserRoleRevocationData,
     UserRoleRevocationInput,
 )
 from ai.backend.manager.data.permission.status import (
@@ -55,7 +61,13 @@
 from ai.backend.manager.models.rbac_models.user_role import UserRoleRow
 from ai.backend.manager.models.user import UserRow
 from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
-from ai.backend.manager.repositories.base.creator import Creator, execute_creator
+from ai.backend.manager.repositories.base.creator import (
+    BulkCreator,
+    BulkCreatorResultWithFailures,
+    Creator,
+    execute_bulk_creator_partial,
+    execute_creator,
+)
 from ai.backend.manager.repositories.base.purger import Purger, execute_purger
 from ai.backend.manager.repositories.base.querier import BatchQuerier, execute_batch_querier
 from ai.backend.manager.repositories.base.updater import Updater, execute_updater
@@ -69,6 +81,8 @@
     PermissionSearchScope,
 )
 
+log = BraceStyleAdapter(logging.getLogger(__spec__.name))
+
 
 @dataclass
 class CreateRoleInput:
@@ -978,3 +992,50 @@ async def check_permission_with_scope_chain(
         async with self._db.begin_readonly_session_read_committed() as db_session:
             result = await db_session.scalar(combined_query)
             return result or False
+
+    async def bulk_assign_role(
+        self, bulk_creator: BulkCreator[UserRoleRow]
+    ) -> BulkCreatorResultWithFailures[UserRoleRow]:
+        async with self._db.begin_session() as db_session:
+            return await execute_bulk_creator_partial(db_session, bulk_creator)
+
+    async def bulk_revoke_role(
+        self, data: BulkUserRoleRevocationInput
+    ) -> BulkRoleRevocationResultData:
+        successes: list[UserRoleRevocationData] = []
+        failures: list[BulkRoleRevocationFailure] = []
+
+        async with self._db.begin_session() as db_session:
+            for user_id in data.user_ids:
+                try:
+                    async with db_session.begin_nested():
+                        stmt = (
+                            sa.select(UserRoleRow)
+                            .where(UserRoleRow.user_id == user_id)
+                            .where(UserRoleRow.role_id == data.role_id)
+                        )
+                        user_role_row = await db_session.scalar(stmt)
+                        if user_role_row is None:
+                            raise RoleNotAssigned(
+                                f"Role {data.role_id} is not assigned to user {user_id}."
+                            )
+                        user_role_id = user_role_row.id
+                        await db_session.delete(user_role_row)
+                        await db_session.flush()
+                        successes.append(
+                            UserRoleRevocationData(
+                                user_role_id=user_role_id,
+                                user_id=user_id,
+                                role_id=data.role_id,
+                            )
+                        )
+                except Exception as e:
+                    log.warning(
+                        "Failed to revoke role {} from user {}: {}",
+                        data.role_id,
+                        user_id,
+                        str(e),
+                    )
+                    failures.append(BulkRoleRevocationFailure(user_id=user_id, message=str(e)))
+
+        return BulkRoleRevocationResultData(successes=successes, failures=failures)

src/ai/backend/manager/repositories/permission_controller/repository.py

@@ -2,6 +2,7 @@
 
 import uuid
 from collections.abc import Mapping
+from typing import cast
 
 from ai.backend.common.data.permission.types import GLOBAL_SCOPE_ID, OperationType
 from ai.backend.common.exception import BackendAIError
@@ -22,6 +23,10 @@
 from ai.backend.manager.data.permission.role import (
     AssignedUserListResult,
     BatchEntityPermissionCheckInput,
+    BulkRoleAssignmentFailure,
+    BulkRoleAssignmentResultData,
+    BulkRoleRevocationResultData,
+    BulkUserRoleRevocationInput,
     RoleData,
     RoleDetailData,
     RoleListResult,
@@ -42,11 +47,13 @@
 from ai.backend.manager.models.rbac_models.permission.object_permission import ObjectPermissionRow
 from ai.backend.manager.models.rbac_models.permission.permission import PermissionRow
 from ai.backend.manager.models.rbac_models.role import RoleRow
+from ai.backend.manager.models.rbac_models.user_role import UserRoleRow
 from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
-from ai.backend.manager.repositories.base.creator import Creator
+from ai.backend.manager.repositories.base.creator import BulkCreator, Creator
 from ai.backend.manager.repositories.base.purger import Purger
 from ai.backend.manager.repositories.base.querier import BatchQuerier
 from ai.backend.manager.repositories.base.updater import Updater
+from ai.backend.manager.repositories.permission_controller.creators import UserRoleCreatorSpec
 from ai.backend.manager.repositories.permission_controller.types import (
     ObjectPermissionSearchScope,
     PermissionSearchScope,
@@ -179,6 +186,29 @@ async def revoke_role(self, data: UserRoleRevocationInput) -> UserRoleRevocation
             user_role_id=user_role_id, user_id=data.user_id, role_id=data.role_id
         )
 
+    @permission_controller_repository_resilience.apply()
+    async def bulk_assign_role(
+        self, bulk_creator: BulkCreator[UserRoleRow]
+    ) -> BulkRoleAssignmentResultData:
+        result = await self._db_source.bulk_assign_role(bulk_creator)
+        failures = [
+            BulkRoleAssignmentFailure(
+                user_id=cast(UserRoleCreatorSpec, error.spec).user_id,
+                message=str(error.exception),
+            )
+            for error in result.errors
+        ]
+        return BulkRoleAssignmentResultData(
+            successes=[row.to_data() for row in result.successes],
+            failures=failures,
+        )
+
+    @permission_controller_repository_resilience.apply()
+    async def bulk_revoke_role(
+        self, data: BulkUserRoleRevocationInput
+    ) -> BulkRoleRevocationResultData:
+        return await self._db_source.bulk_revoke_role(data)
+
     @permission_controller_repository_resilience.apply()
     async def get_role(self, role_id: uuid.UUID) -> RoleData | None:
         result = await self._db_source.get_role(role_id)

src/ai/backend/manager/services/permission_contoller/actions/__init__.py

@@ -1,4 +1,6 @@
 from .assign_role import AssignRoleAction, AssignRoleActionResult
+from .bulk_assign_role import BulkAssignRoleAction, BulkAssignRoleActionResult
+from .bulk_revoke_role import BulkRevokeRoleAction, BulkRevokeRoleActionResult
 from .check_permission import CheckPermissionAction, CheckPermissionActionResult
 from .create_role import CreateRoleAction, CreateRoleActionResult
 from .delete_role import DeleteRoleAction, DeleteRoleActionResult
@@ -27,6 +29,10 @@
 __all__ = [
     "AssignRoleAction",
     "AssignRoleActionResult",
+    "BulkAssignRoleAction",
+    "BulkAssignRoleActionResult",
+    "BulkRevokeRoleAction",
+    "BulkRevokeRoleActionResult",
     "CheckPermissionAction",
     "CheckPermissionActionResult",
     "CreateRoleAction",

src/ai/backend/manager/services/permission_contoller/actions/bulk_assign_role.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from typing import override
+
+from ai.backend.common.data.permission.types import EntityType
+from ai.backend.manager.actions.action import BaseActionResult
+from ai.backend.manager.actions.types import ActionOperationType
+from ai.backend.manager.data.permission.role import (
+    BulkRoleAssignmentResultData,
+)
+from ai.backend.manager.models.rbac_models.user_role import UserRoleRow
+from ai.backend.manager.repositories.base.creator import BulkCreator
+from ai.backend.manager.services.permission_contoller.actions.base import RoleAction
+
+
+@dataclass
+class BulkAssignRoleAction(RoleAction):
+    bulk_creator: BulkCreator[UserRoleRow]
+
+    @override
+    @classmethod
+    def entity_type(cls) -> EntityType:
+        return EntityType.ROLE_ASSIGNMENT
+
+    @override
+    def entity_id(self) -> str | None:
+        return None
+
+    @override
+    @classmethod
+    def operation_type(cls) -> ActionOperationType:
+        return ActionOperationType.CREATE
+
+
+@dataclass
+class BulkAssignRoleActionResult(BaseActionResult):
+    data: BulkRoleAssignmentResultData
+
+    @override
+    def entity_id(self) -> str | None:
+        return None

src/ai/backend/manager/services/permission_contoller/actions/bulk_revoke_role.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass
+from typing import override
+
+from ai.backend.common.data.permission.types import EntityType
+from ai.backend.manager.actions.action import BaseActionResult
+from ai.backend.manager.actions.types import ActionOperationType
+from ai.backend.manager.data.permission.role import (
+    BulkRoleRevocationResultData,
+    BulkUserRoleRevocationInput,
+)
+from ai.backend.manager.services.permission_contoller.actions.base import RoleAction
+
+
+@dataclass
+class BulkRevokeRoleAction(RoleAction):
+    input: BulkUserRoleRevocationInput
+
+    @override
+    @classmethod
+    def entity_type(cls) -> EntityType:
+        return EntityType.ROLE_ASSIGNMENT
+
+    @override
+    def entity_id(self) -> str | None:
+        return None
+
+    @override
+    @classmethod
+    def operation_type(cls) -> ActionOperationType:
+        return ActionOperationType.DELETE
+
+
+@dataclass
+class BulkRevokeRoleActionResult(BaseActionResult):
+    data: BulkRoleRevocationResultData
+
+    @override
+    def entity_id(self) -> str | None:
+        return None

src/ai/backend/manager/services/permission_contoller/processors.py

@@ -7,6 +7,10 @@
 from .actions import (
     AssignRoleAction,
     AssignRoleActionResult,
+    BulkAssignRoleAction,
+    BulkAssignRoleActionResult,
+    BulkRevokeRoleAction,
+    BulkRevokeRoleActionResult,
     CreateRoleAction,
     CreateRoleActionResult,
     DeleteRoleAction,
@@ -66,6 +70,8 @@ class PermissionControllerProcessors(AbstractProcessorPackage):
     delete_role: ActionProcessor[DeleteRoleAction, DeleteRoleActionResult]
     assign_role: ActionProcessor[AssignRoleAction, AssignRoleActionResult]
     revoke_role: ActionProcessor[RevokeRoleAction, RevokeRoleActionResult]
+    bulk_assign_role: ActionProcessor[BulkAssignRoleAction, BulkAssignRoleActionResult]
+    bulk_revoke_role: ActionProcessor[BulkRevokeRoleAction, BulkRevokeRoleActionResult]
     get_role_detail: ActionProcessor[GetRoleDetailAction, GetRoleDetailActionResult]
     search_roles: ActionProcessor[SearchRolesAction, SearchRolesActionResult]
     search_users_assigned_to_role: ActionProcessor[
@@ -94,6 +100,8 @@ def __init__(
         self.purge_role = ActionProcessor(service.purge_role, action_monitors)
         self.assign_role = ActionProcessor(service.assign_role, action_monitors)
         self.revoke_role = ActionProcessor(service.revoke_role, action_monitors)
+        self.bulk_assign_role = ActionProcessor(service.bulk_assign_role, action_monitors)
+        self.bulk_revoke_role = ActionProcessor(service.bulk_revoke_role, action_monitors)
         self.get_role_detail = ActionProcessor(service.get_role_detail, action_monitors)
         self.search_roles = ActionProcessor(service.search_roles, action_monitors)
         self.search_users_assigned_to_role = ActionProcessor(
@@ -122,6 +130,8 @@ def supported_actions(self) -> list[ActionSpec]:
             PurgeRoleAction.spec(),
             AssignRoleAction.spec(),
             RevokeRoleAction.spec(),
+            BulkAssignRoleAction.spec(),
+            BulkRevokeRoleAction.spec(),
             GetRoleDetailAction.spec(),
             SearchRolesAction.spec(),
             SearchUsersAssignedToRoleAction.spec(),

src/ai/backend/manager/services/permission_contoller/service.py

@@ -12,6 +12,14 @@
     AssignRoleAction,
     AssignRoleActionResult,
 )
+from ai.backend.manager.services.permission_contoller.actions.bulk_assign_role import (
+    BulkAssignRoleAction,
+    BulkAssignRoleActionResult,
+)
+from ai.backend.manager.services.permission_contoller.actions.bulk_revoke_role import (
+    BulkRevokeRoleAction,
+    BulkRevokeRoleActionResult,
+)
 from ai.backend.manager.services.permission_contoller.actions.create_role import (
     CreateRoleAction,
     CreateRoleActionResult,
@@ -186,6 +194,16 @@ async def revoke_role(self, action: RevokeRoleAction) -> RevokeRoleActionResult:
 
         return RevokeRoleActionResult(data=data)
 
+    async def bulk_assign_role(self, action: BulkAssignRoleAction) -> BulkAssignRoleActionResult:
+        """Assigns a role to multiple users with partial failure support."""
+        data = await self._repository.bulk_assign_role(action.bulk_creator)
+        return BulkAssignRoleActionResult(data=data)
+
+    async def bulk_revoke_role(self, action: BulkRevokeRoleAction) -> BulkRevokeRoleActionResult:
+        """Revokes a role from multiple users with partial failure support."""
+        data = await self._repository.bulk_revoke_role(action.input)
+        return BulkRevokeRoleActionResult(data=data)
+
     async def get_role_detail(self, action: GetRoleDetailAction) -> GetRoleDetailActionResult:
         """Get role with all permission details and assigned users."""
         role_data = await self._repository.get_role_with_permissions(action.role_id)