feat(BA-6254): add superadmin ensure-system-role repository/service/action

changes/11891.feature.md

@@ -0,0 +1 @@
+Add a superadmin operation that ensures a user has an RBAC system role, creating it (with permissions and the user-role mapping) when missing.

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

@@ -39,3 +39,9 @@ class UserRoleDataWithRole:
     deleted_at: datetime | None
 
     mapped_role_data: RoleData
+
+
+@dataclass(frozen=True)
+class RoleMappingData:
+    user_id: uuid.UUID
+    role_data: RoleData

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

@@ -13,6 +13,7 @@
 from ai.backend.common.data.permission.types import (
     RBACElementType,
     RelationType,
+    RoleSource,
 )
 from ai.backend.logging.utils import BraceStyleAdapter
 from ai.backend.manager.actions.action.rbac_role_invitation import (
@@ -39,6 +40,7 @@
     BulkUserRoleRevocationInput,
     PermissionResolutionKey,
     ProjectRoleCount,
+    RoleData,
     RoleListResult,
     RolePermissionsUpdateInput,
     RoleRevocationResult,
@@ -58,6 +60,7 @@
     ScopeListResult,
     ScopeType,
 )
+from ai.backend.manager.data.permission.user_role import RoleMappingData
 from ai.backend.manager.data.role_invitation.types import (
     RoleInvitationData,
     RoleInvitationState,
@@ -106,6 +109,10 @@
     PermissionCreatorSpec,
     UserRoleCreatorSpec,
 )
+from ai.backend.manager.repositories.permission_controller.role_manager import (
+    RoleManager,
+    UserSystemRoleSpec,
+)
 from ai.backend.manager.repositories.permission_controller.types import (
     PermissionSearchScope,
     ScopedRoleSearchScope,
@@ -148,9 +155,56 @@ class _PermissionGroupKey:
 
 class PermissionDBSource:
     _db: ExtendedAsyncSAEngine
+    _role_manager: RoleManager
 
     def __init__(self, db: ExtendedAsyncSAEngine) -> None:
         self._db = db
+        self._role_manager = RoleManager()
+
+    async def ensure_user_system_roles(
+        self, specs: Collection[UserSystemRoleSpec]
+    ) -> list[RoleMappingData]:
+        """Ensure the SYSTEM role(s) for the given scope exist.
+
+        Idempotent: roles already present in the scope (matched by name) are
+        reused; missing ones are created together with their permissions. For
+        the USER scope, the user-to-role mapping is ensured as well.
+        """
+        if not specs:
+            return []
+        async with self._db.begin_session_read_committed() as db_session:
+            existing = await self._existing_user_system_roles(db_session, specs)
+            results: list[RoleMappingData] = []
+            for spec in specs:
+                role = existing.get(spec.user_id)
+                if role is None:
+                    # The user has no SYSTEM role yet: create it and map the
+                    # user to it. Existing mappings are left untouched so the
+                    # operation stays idempotent on repeated calls.
+                    role = await self._role_manager.create_system_role(db_session, spec)
+                    await execute_creator(
+                        db_session,
+                        Creator(spec=UserRoleCreatorSpec(user_id=spec.user_id, role_id=role.id)),
+                    )
+                results.append(RoleMappingData(user_id=spec.user_id, role_data=role))
+            return results
+
+    async def _existing_user_system_roles(
+        self, db_session: SASession, specs: Iterable[UserSystemRoleSpec]
+    ) -> dict[uuid.UUID, RoleData]:
+        stmt = (
+            sa.select(UserRoleRow.user_id, RoleRow)
+            .select_from(UserRoleRow)
+            .join(RoleRow, UserRoleRow.role_id == RoleRow.id)
+            .where(
+                sa.and_(
+                    UserRoleRow.user_id.in_([spec.user_id for spec in specs]),
+                    RoleRow.source == RoleSource.SYSTEM,
+                )
+            )
+        )
+        result = await db_session.execute(stmt)
+        return {row.user_id: row.RoleRow.to_data() for row in result}
 
     @staticmethod
     async def _sync_user_scopes_on_assign(

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

@@ -50,6 +50,7 @@
     ScopeListResult,
     ScopeType,
 )
+from ai.backend.manager.data.permission.user_role import RoleMappingData
 from ai.backend.manager.data.role_invitation.types import RoleInvitationData
 from ai.backend.manager.models.rbac_models.permission.permission import PermissionRow
 from ai.backend.manager.models.rbac_models.role import RoleRow
@@ -66,6 +67,7 @@
     PermissionCreatorSpec,
     UserRoleCreatorSpec,
 )
+from ai.backend.manager.repositories.permission_controller.role_manager import UserSystemRoleSpec
 from ai.backend.manager.repositories.permission_controller.types import (
     PermissionSearchScope,
     ScopedRoleSearchScope,
@@ -114,6 +116,17 @@ async def create_role(self, input_data: CreateRoleInput) -> RoleData:
         role_row = await self._db_source.create_role(input_data)
         return role_row.to_data()
 
+    @permission_controller_repository_resilience.apply()
+    async def ensure_user_system_roles(
+        self, specs: Collection[UserSystemRoleSpec]
+    ) -> list[RoleMappingData]:
+        """
+        Ensure the SYSTEM role(s) for the given scope exist (idempotent).
+
+        Returns the ensured roles.
+        """
+        return await self._db_source.ensure_user_system_roles(specs)
+
     @permission_controller_repository_resilience.apply()
     async def create_permission(
         self,

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

@@ -12,6 +12,7 @@
 from .check_permission import CheckPermissionAction, CheckPermissionActionResult
 from .create_role import CreateRoleAction, CreateRoleActionResult
 from .delete_role import DeleteRoleAction, DeleteRoleActionResult
+from .ensure_system_role import EnsureSystemRoleAction, EnsureSystemRoleActionResult
 from .get_permission_matrix import GetPermissionMatrixAction, GetPermissionMatrixActionResult
 from .get_role_detail import GetRoleDetailAction, GetRoleDetailActionResult
 from .purge_role import PurgeRoleAction, PurgeRoleActionResult
@@ -61,6 +62,8 @@
     "CreateRoleActionResult",
     "DeleteRoleAction",
     "DeleteRoleActionResult",
+    "EnsureSystemRoleAction",
+    "EnsureSystemRoleActionResult",
     "GetPermissionMatrixAction",
     "GetPermissionMatrixActionResult",
     "GetRoleDetailAction",

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

@@ -0,0 +1,34 @@
+from collections.abc import Collection
+from dataclasses import dataclass
+from typing import override
+
+from ai.backend.manager.actions.action import BaseActionResult
+from ai.backend.manager.actions.types import ActionOperationType
+from ai.backend.manager.data.permission.user_role import RoleMappingData
+from ai.backend.manager.repositories.permission_controller.role_manager import UserSystemRoleSpec
+from ai.backend.manager.services.permission_contoller.actions.base import RoleAction
+
+
+@dataclass
+class EnsureSystemRoleAction(RoleAction):
+    """Ensure the SYSTEM role(s) for the given users exist (superadmin only)."""
+
+    specs: Collection[UserSystemRoleSpec]
+
+    @override
+    def entity_id(self) -> str | None:
+        return None
+
+    @override
+    @classmethod
+    def operation_type(cls) -> ActionOperationType:
+        return ActionOperationType.CREATE
+
+
+@dataclass
+class EnsureSystemRoleActionResult(BaseActionResult):
+    data: list[RoleMappingData]
+
+    @override
+    def entity_id(self) -> str | None:
+        return None

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

@@ -22,6 +22,8 @@
     CreateRoleActionResult,
     DeleteRoleAction,
     DeleteRoleActionResult,
+    EnsureSystemRoleAction,
+    EnsureSystemRoleActionResult,
     GetRoleDetailAction,
     GetRoleDetailActionResult,
     PurgeRoleAction,
@@ -107,6 +109,7 @@ class PermissionControllerProcessors(AbstractProcessorPackage):
     """Processor package for RBAC permission controller operations."""
 
     create_role: ActionProcessor[CreateRoleAction, CreateRoleActionResult]
+    ensure_system_role: ActionProcessor[EnsureSystemRoleAction, EnsureSystemRoleActionResult]
     update_role: ActionProcessor[UpdateRoleAction, UpdateRoleActionResult]
     delete_role: ActionProcessor[DeleteRoleAction, DeleteRoleActionResult]
     assign_role: ActionProcessor[AssignRoleAction, AssignRoleActionResult]
@@ -179,6 +182,7 @@ def __init__(
         validators: ActionValidators,
     ) -> None:
         self.create_role = ActionProcessor(service.create_role, action_monitors)
+        self.ensure_system_role = ActionProcessor(service.ensure_system_role, action_monitors)
         self.update_role = ActionProcessor(service.update_role, action_monitors)
         self.delete_role = ActionProcessor(service.delete_role, action_monitors)
         self.purge_role = ActionProcessor(service.purge_role, action_monitors)
@@ -256,6 +260,7 @@ def __init__(
     def supported_actions(self) -> list[ActionSpec]:
         return [
             CreateRoleAction.spec(),
+            EnsureSystemRoleAction.spec(),
             UpdateRoleAction.spec(),
             DeleteRoleAction.spec(),
             PurgeRoleAction.spec(),

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

@@ -49,6 +49,10 @@
     DeleteRoleAction,
     DeleteRoleActionResult,
 )
+from ai.backend.manager.services.permission_contoller.actions.ensure_system_role import (
+    EnsureSystemRoleAction,
+    EnsureSystemRoleActionResult,
+)
 from ai.backend.manager.services.permission_contoller.actions.get_entity_types import (
     GetEntityTypesAction,
     GetEntityTypesActionResult,
@@ -196,6 +200,15 @@ async def create_role(self, action: CreateRoleAction) -> CreateRoleActionResult:
             data=result,
         )
 
+    async def ensure_system_role(
+        self, action: EnsureSystemRoleAction
+    ) -> EnsureSystemRoleActionResult:
+        """
+        Ensures the SYSTEM role(s) for the given users exist (idempotent).
+        """
+        roles = await self._repository.ensure_user_system_roles(action.specs)
+        return EnsureSystemRoleActionResult(data=roles)
+
     async def create_permission(
         self, action: CreatePermissionAction
     ) -> CreatePermissionActionResult: