feat(BA-4526): Unify ScopeType/EntityType to RBACElementType in domain repos (#10336)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: fregataa

changes/10336.enhance.md

@@ -0,0 +1 @@
+Unify ScopeType/EntityType to RBACElementType in domain data types, non-permission-controller repositories, and vfolder grant/revoke operations

src/ai/backend/manager/data/domain/types.py

@@ -6,6 +6,7 @@
 from datetime import datetime
 from typing import Any, override
 
+from ai.backend.common.data.permission.types import RBACElementType
 from ai.backend.common.data.user.types import UserRole
 from ai.backend.common.types import ResourceSlot, VFolderHostPermissionMap
 from ai.backend.manager.data.permission.id import ScopeId
@@ -46,9 +47,9 @@ def scope_id(self) -> ScopeId:
     def role_name(self) -> str:
         return f"domain-{self.name}-admin"
 
-    def entity_operations(self) -> Mapping[EntityType, Iterable[OperationType]]:
+    def entity_operations(self) -> Mapping[RBACElementType, Iterable[OperationType]]:
         return {
-            entity: OperationType.admin_operations()
+            entity.to_element(): OperationType.admin_operations()
             for entity in EntityType.admin_accessible_entity_types_in_domain()
         }
 

src/ai/backend/manager/data/group/types.py

@@ -7,6 +7,7 @@
 from datetime import datetime
 from typing import Any, override
 
+from ai.backend.common.data.permission.types import RBACElementType
 from ai.backend.common.types import ResourceSlot, VFolderHostPermissionMap
 from ai.backend.manager.data.permission.id import ScopeId
 from ai.backend.manager.data.permission.types import (
@@ -63,9 +64,9 @@ def scope_id(self) -> ScopeId:
     def role_name(self) -> str:
         return f"project-{str(self.id)[:8]}-admin"
 
-    def entity_operations(self) -> Mapping[EntityType, Iterable[OperationType]]:
+    def entity_operations(self) -> Mapping[RBACElementType, Iterable[OperationType]]:
         return {
-            entity: OperationType.admin_operations()
+            entity.to_element(): OperationType.admin_operations()
             for entity in EntityType.admin_accessible_entity_types_in_project()
         }
 

src/ai/backend/manager/data/user/types.py

@@ -9,6 +9,7 @@
 
 from sqlalchemy.engine import Row
 
+from ai.backend.common.data.permission.types import RBACElementType
 from ai.backend.common.data.user.types import UserRole
 from ai.backend.common.types import AccessKey
 from ai.backend.manager.data.keypair.types import KeyPairData
@@ -97,13 +98,13 @@ def scope_id(self) -> ScopeId:
     def role_name(self) -> str:
         return f"user-{str(self.id)[:8]}"
 
-    def entity_operations(self) -> Mapping[EntityType, Iterable[OperationType]]:
+    def entity_operations(self) -> Mapping[RBACElementType, Iterable[OperationType]]:
         resource_entity_permissions = {
-            entity: OperationType.owner_operations()
+            entity.to_element(): OperationType.owner_operations()
             for entity in EntityType.owner_accessible_entity_types_in_user()
         }
         user_permissions = OperationType.owner_operations() - {OperationType.CREATE}
-        return {EntityType.USER: user_permissions, **resource_entity_permissions}
+        return {RBACElementType.USER: user_permissions, **resource_entity_permissions}
 
     @classmethod
     def from_row(cls, row: Row[Any]) -> Self:

src/ai/backend/manager/repositories/base/rbac/granter.py

@@ -8,8 +8,8 @@
 
 from ai.backend.common.data.permission.types import (
     OperationType,
+    RBACElementType,
     RelationType,
-    ScopeType,
 )
 from ai.backend.manager.data.permission.id import (
     ObjectId,
@@ -45,7 +45,7 @@ class RBACGranter:
     """
 
     granted_entity_id: ObjectId
-    granted_entity_scope_type: ScopeType
+    granted_entity_scope_type: RBACElementType
     target_scope_id: ScopeId
     target_role_ids: list[UUID]
     operations: list[OperationType]
@@ -97,7 +97,7 @@ async def execute_rbac_granter(
     perms = [
         PermissionRow(
             role_id=role_id,
-            scope_type=granter.granted_entity_scope_type,
+            scope_type=granter.granted_entity_scope_type.to_scope_type(),
             scope_id=entity_id.entity_id,
             entity_type=entity_id.entity_type,
             operation=operation,

src/ai/backend/manager/repositories/base/rbac/revoker.py

@@ -9,7 +9,7 @@
 from sqlalchemy.engine import CursorResult
 from sqlalchemy.ext.asyncio import AsyncSession as SASession
 
-from ai.backend.common.data.permission.types import OperationType, ScopeType
+from ai.backend.common.data.permission.types import OperationType, RBACElementType
 from ai.backend.manager.data.permission.id import ObjectId
 from ai.backend.manager.models.rbac_models.permission.permission import PermissionRow
 
@@ -36,7 +36,7 @@ class RBACRevoker:
     """
 
     entity_id: ObjectId
-    entity_scope_type: ScopeType
+    entity_scope_type: RBACElementType
     target_role_ids: list[UUID]
     operations: list[OperationType] | None = None
 
@@ -50,7 +50,7 @@ async def _delete_permissions(
     db_sess: SASession,
     role_ids: Collection[UUID],
     entity_id: ObjectId,
-    scope_type: ScopeType,
+    scope_type: RBACElementType,
     operations: list[OperationType] | None,
 ) -> int:
     """Delete permissions for the given entity-as-scope and roles."""
@@ -59,7 +59,7 @@ async def _delete_permissions(
 
     conditions = [
         PermissionRow.role_id.in_(role_ids),
-        PermissionRow.scope_type == scope_type,
+        PermissionRow.scope_type == scope_type.to_scope_type(),
         PermissionRow.scope_id == entity_id.entity_id,
         PermissionRow.entity_type == entity_id.entity_type,
     ]

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

@@ -6,6 +6,7 @@
 import sqlalchemy as sa
 from sqlalchemy.ext.asyncio import AsyncSession as SASession
 
+from ai.backend.common.data.permission.types import RBACElementType
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.manager.data.permission.id import ObjectId, ScopeId
 from ai.backend.manager.data.permission.object_permission import ObjectPermissionData
@@ -15,7 +16,6 @@
 )
 from ai.backend.manager.data.permission.status import RoleStatus
 from ai.backend.manager.data.permission.types import (
-    EntityType,
     OperationType,
     RoleSource,
 )
@@ -42,7 +42,7 @@ def scope_id(self) -> ScopeId: ...
 
     def role_name(self) -> str: ...
 
-    def entity_operations(self) -> Mapping[EntityType, Iterable[OperationType]]:
+    def entity_operations(self) -> Mapping[RBACElementType, Iterable[OperationType]]:
         """Returns a mapping of entity types to the set of operations that should be granted for each entity type."""
         ...
 
@@ -78,13 +78,13 @@ async def _create_permissions(
         role_id: uuid.UUID,
     ) -> list[PermissionData]:
         permission_rows: list[PermissionRow] = []
-        for entity, operations in data.entity_operations().items():
+        for element_type, operations in data.entity_operations().items():
             for operation in operations:
                 creator = PermissionCreator(
                     role_id=role_id,
                     scope_type=data.scope_id().scope_type,
                     scope_id=data.scope_id().scope_id,
-                    entity_type=entity,
+                    entity_type=element_type.to_entity_type(),
                     operation=operation,
                 )
                 permission_rows.append(PermissionRow.from_input(creator))

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

@@ -349,7 +349,7 @@ async def create_vfolder_with_permission(
                         entity_type=EntityType.VFOLDER,
                         entity_id=str(params.id),
                     ),
-                    granted_entity_scope_type=ScopeType.VFOLDER,
+                    granted_entity_scope_type=RBACElementType.VFOLDER,
                     target_scope_id=ScopeId(
                         scope_type=ScopeType.USER,
                         scope_id=str(params.user),
@@ -538,7 +538,7 @@ async def create_vfolder_permission(
                     entity_type=EntityType.VFOLDER,
                     entity_id=str(vfolder_id),
                 ),
-                granted_entity_scope_type=ScopeType.VFOLDER,
+                granted_entity_scope_type=RBACElementType.VFOLDER,
                 target_scope_id=ScopeId(
                     scope_type=ScopeType.USER,
                     scope_id=str(user_id),
@@ -577,7 +577,7 @@ async def delete_vfolder_permission(self, vfolder_id: uuid.UUID, user_id: uuid.U
                     entity_type=EntityType.VFOLDER,
                     entity_id=str(vfolder_id),
                 ),
-                entity_scope_type=ScopeType.VFOLDER,
+                entity_scope_type=RBACElementType.VFOLDER,
                 target_role_ids=[user_role_id],
                 operations=None,  # Revoke all operations
             )

tests/unit/manager/repositories/base/rbac/test_granter.py

@@ -11,7 +11,7 @@
 import pytest
 import sqlalchemy as sa
 
-from ai.backend.common.data.permission.types import OperationType, RelationType
+from ai.backend.common.data.permission.types import OperationType, RBACElementType, RelationType
 from ai.backend.manager.data.permission.id import ObjectId, ScopeId
 from ai.backend.manager.data.permission.types import (
     EntityType,
@@ -53,7 +53,7 @@
 class GranterTestContext:
     """Context data for granter tests."""
 
-    entity_scope_type: ScopeType
+    entity_scope_type: RBACElementType
     entity_id: ObjectId
     target_scope_id: ScopeId
 
@@ -116,7 +116,7 @@ async def single_role(
             role_id = role.id
 
         yield SingleRoleContext(
-            entity_scope_type=ScopeType.VFOLDER,
+            entity_scope_type=RBACElementType.VFOLDER,
             entity_id=entity_id,
             target_scope_id=target_scope_id,
             role_id=role_id,
@@ -133,7 +133,7 @@ async def empty_context(
         target_scope_id = ScopeId(scope_type=ScopeType.USER, scope_id=str(uuid.uuid4()))
 
         yield GranterTestContext(
-            entity_scope_type=ScopeType.VFOLDER,
+            entity_scope_type=RBACElementType.VFOLDER,
             entity_id=entity_id,
             target_scope_id=target_scope_id,
         )
@@ -232,7 +232,7 @@ async def multi_role_context(
                 role_ids.append(role.id)
 
         yield MultiRoleContext(
-            entity_scope_type=ScopeType.VFOLDER,
+            entity_scope_type=RBACElementType.VFOLDER,
             entity_id=entity_id,
             target_scope_id=target_scope_id,
             role_ids=role_ids,
@@ -296,7 +296,7 @@ async def single_role(
             role_id = role.id
 
         yield SingleRoleContext(
-            entity_scope_type=ScopeType.VFOLDER,
+            entity_scope_type=RBACElementType.VFOLDER,
             entity_id=entity_id,
             target_scope_id=target_scope_id,
             role_id=role_id,

tests/unit/manager/repositories/base/rbac/test_revoker.py

@@ -11,7 +11,11 @@
 import pytest
 import sqlalchemy as sa
 
-from ai.backend.common.data.permission.types import EntityType, OperationType, ScopeType
+from ai.backend.common.data.permission.types import (
+    EntityType,
+    OperationType,
+    RBACElementType,
+)
 from ai.backend.manager.data.permission.id import ObjectId
 from ai.backend.manager.data.permission.types import RoleSource
 from ai.backend.manager.models.rbac_models.permission.permission import PermissionRow
@@ -46,7 +50,7 @@ class SingleEntityWithRoleContext:
     """Context with single entity granted to a role."""
 
     entity_id: ObjectId
-    entity_scope_type: ScopeType
+    entity_scope_type: RBACElementType
     role_id: UUID
 
 
@@ -55,7 +59,7 @@ class EntityWithTwoRolesContext:
     """Context with entity granted to two different roles."""
 
     entity_id: ObjectId
-    entity_scope_type: ScopeType
+    entity_scope_type: RBACElementType
     role_id1: UUID
     role_id2: UUID
 
@@ -90,7 +94,7 @@ async def single_entity_with_role(
     ) -> AsyncGenerator[SingleEntityWithRoleContext, None]:
         """Create entity with role having permissions."""
         entity_id = ObjectId(entity_type=EntityType.VFOLDER, entity_id=str(uuid.uuid4()))
-        entity_scope_type = ScopeType.VFOLDER
+        entity_scope_type = RBACElementType.VFOLDER
 
         role_id: UUID
 
@@ -108,7 +112,7 @@ async def single_entity_with_role(
             for op in [OperationType.READ, OperationType.UPDATE]:
                 perm = PermissionRow(
                     role_id=role.id,
-                    scope_type=entity_scope_type,
+                    scope_type=entity_scope_type.to_scope_type(),
                     scope_id=entity_id.entity_id,
                     entity_type=entity_id.entity_type,
                     operation=op,
@@ -231,7 +235,7 @@ async def entity_with_two_roles(
     ) -> AsyncGenerator[EntityWithTwoRolesContext, None]:
         """Create entity granted to two different roles."""
         entity_id = ObjectId(entity_type=EntityType.VFOLDER, entity_id=str(uuid.uuid4()))
-        entity_scope_type = ScopeType.VFOLDER
+        entity_scope_type = RBACElementType.VFOLDER
 
         role_id1: UUID
         role_id2: UUID
@@ -256,7 +260,7 @@ async def entity_with_two_roles(
             for role in [role1, role2]:
                 perm = PermissionRow(
                     role_id=role.id,
-                    scope_type=entity_scope_type,
+                    scope_type=entity_scope_type.to_scope_type(),
                     scope_id=entity_id.entity_id,
                     entity_type=entity_id.entity_type,
                     operation=OperationType.READ,
@@ -331,7 +335,7 @@ async def single_entity_with_role(
     ) -> AsyncGenerator[SingleEntityWithRoleContext, None]:
         """Create entity with role having permissions."""
         entity_id = ObjectId(entity_type=EntityType.VFOLDER, entity_id=str(uuid.uuid4()))
-        entity_scope_type = ScopeType.VFOLDER
+        entity_scope_type = RBACElementType.VFOLDER
 
         role_id: UUID
 
@@ -346,7 +350,7 @@ async def single_entity_with_role(
 
             perm = PermissionRow(
                 role_id=role.id,
-                scope_type=entity_scope_type,
+                scope_type=entity_scope_type.to_scope_type(),
                 scope_id=entity_id.entity_id,
                 entity_type=entity_id.entity_type,
                 operation=OperationType.READ,