diff --git a/changes/11499.feature.md b/changes/11499.feature.md new file mode 100644 index 00000000000..b71cb5b544a --- /dev/null +++ b/changes/11499.feature.md @@ -0,0 +1 @@ +Add a request-scoped GraphQL DataLoader for batched per-entity effective permission resolution. diff --git a/src/ai/backend/manager/api/adapters/rbac/adapter.py b/src/ai/backend/manager/api/adapters/rbac/adapter.py index be70202fb36..fba2d7de3ae 100644 --- a/src/ai/backend/manager/api/adapters/rbac/adapter.py +++ b/src/ai/backend/manager/api/adapters/rbac/adapter.py @@ -4,7 +4,7 @@ import uuid from collections import defaultdict -from collections.abc import Callable, Collection, Sequence +from collections.abc import Callable, Collection, Mapping, Sequence from datetime import UTC, datetime from functools import lru_cache from uuid import UUID @@ -175,6 +175,7 @@ BulkRolePermissionReplaceResultData, BulkRoleRevocationResultData, BulkUserRoleRevocationInput, + PermissionResolutionKey, RoleData, RoleDetailData, UserRoleAssignmentData, @@ -268,6 +269,9 @@ from ai.backend.manager.services.permission_contoller.actions.replace_role_permissions import ( ReplaceRolePermissionsAction, ) +from ai.backend.manager.services.permission_contoller.actions.resolve_effective_permissions import ( + ResolveEffectivePermissionsAction, +) from ai.backend.manager.services.permission_contoller.actions.revoke_role import RevokeRoleAction from ai.backend.manager.services.permission_contoller.actions.search_element_associations import ( SearchElementAssociationsAction, @@ -635,6 +639,18 @@ async def get_permission_matrix(self) -> list[ScopeEntityOperationCombinationInf for scope, entity_map in sorted(matrix.items(), key=lambda e: e[0].value) ] + # ------------------------------------------------------------------ effective permissions + + async def batch_resolve_effective_permissions( + self, + keys: Sequence[PermissionResolutionKey], + ) -> Mapping[PermissionResolutionKey, frozenset[InternalOperationType]]: + """Resolve granted operations for each input key; missing keys map to an empty frozenset.""" + action_result = await self._processors.permission_controller.resolve_effective_permissions.wait_for_complete( + ResolveEffectivePermissionsAction(keys=list(keys)) + ) + return action_result.permissions + # ------------------------------------------------------------------ create async def create(self, input: CreateRoleInput) -> CreateRolePayload: diff --git a/src/ai/backend/manager/api/gql/data_loader/data_loaders.py b/src/ai/backend/manager/api/gql/data_loader/data_loaders.py index 662fc1e631f..8df74336c0b 100644 --- a/src/ai/backend/manager/api/gql/data_loader/data_loaders.py +++ b/src/ai/backend/manager/api/gql/data_loader/data_loaders.py @@ -6,9 +6,11 @@ from strawberry.dataloader import DataLoader +from ai.backend.common.data.permission.types import OperationType from ai.backend.common.identifier.deployment import DeploymentID from ai.backend.common.types import AgentId, ImageID, KernelId, SessionId from ai.backend.manager.data.permission.id import ObjectId +from ai.backend.manager.data.permission.role import PermissionResolutionKey if TYPE_CHECKING: from ai.backend.common.dto.manager.v2.rbac.response import EntityNode # pants: no-infer-dep @@ -684,6 +686,20 @@ async def load_fn(ids: list[uuid.UUID]) -> list[RoleGQL | None]: return DataLoader(load_fn=load_fn) + @cached_property + def effective_permissions_loader( + self, + ) -> DataLoader[PermissionResolutionKey, frozenset[OperationType]]: + adapter = self._adapters.rbac + + async def load_fn( + keys: list[PermissionResolutionKey], + ) -> list[frozenset[OperationType]]: + result = await adapter.batch_resolve_effective_permissions(keys) + return [result.get(key, frozenset()) for key in keys] + + return DataLoader(load_fn=load_fn) + @cached_property def permission_loader( self, diff --git a/src/ai/backend/manager/services/permission_contoller/processors.py b/src/ai/backend/manager/services/permission_contoller/processors.py index d458bce1e83..115b9d7473e 100644 --- a/src/ai/backend/manager/services/permission_contoller/processors.py +++ b/src/ai/backend/manager/services/permission_contoller/processors.py @@ -27,6 +27,8 @@ PurgeRoleAction, ReplaceRolePermissionsAction, ReplaceRolePermissionsActionResult, + ResolveEffectivePermissionsAction, + ResolveEffectivePermissionsActionResult, RevokeRoleAction, RevokeRoleActionResult, SearchRolesAction, @@ -139,6 +141,9 @@ class PermissionControllerProcessors(AbstractProcessorPackage): get_permission_matrix: ActionProcessor[ GetPermissionMatrixAction, GetPermissionMatrixActionResult ] + resolve_effective_permissions: ActionProcessor[ + ResolveEffectivePermissionsAction, ResolveEffectivePermissionsActionResult + ] search_entities: ActionProcessor[SearchEntitiesAction, SearchEntitiesActionResult] search_element_associations: ActionProcessor[ SearchElementAssociationsAction, SearchElementAssociationsActionResult @@ -211,6 +216,9 @@ def __init__( self.get_scope_types = ActionProcessor(service.get_scope_types, action_monitors) self.get_entity_types = ActionProcessor(service.get_entity_types, action_monitors) self.get_permission_matrix = ActionProcessor(service.get_permission_matrix, action_monitors) + self.resolve_effective_permissions = ActionProcessor( + service.resolve_effective_permissions, action_monitors + ) self.search_entities = ActionProcessor(service.search_entities, action_monitors) self.search_element_associations = ActionProcessor( service.search_element_associations, action_monitors @@ -275,6 +283,7 @@ def supported_actions(self) -> list[ActionSpec]: GetScopeTypesAction.spec(), GetEntityTypesAction.spec(), GetPermissionMatrixAction.spec(), + ResolveEffectivePermissionsAction.spec(), SearchEntitiesAction.spec(), SearchElementAssociationsAction.spec(), SearchPermissionsAction.spec(),