diff --git a/changes/11179.feature.md b/changes/11179.feature.md new file mode 100644 index 00000000000..562e467360c --- /dev/null +++ b/changes/11179.feature.md @@ -0,0 +1 @@ +Add my-role REST v2 SDK endpoint and CLI command for users to query their own role assignments diff --git a/docs/manager/graphql-reference/supergraph.graphql b/docs/manager/graphql-reference/supergraph.graphql index 700055bf40d..e2157053046 100644 --- a/docs/manager/graphql-reference/supergraph.graphql +++ b/docs/manager/graphql-reference/supergraph.graphql @@ -14857,7 +14857,7 @@ type RoleAssignmentEdge input RoleAssignmentFilter @join__type(graph: STRAWBERRY) { - roleId: UUID = null + roleId: UUIDFilter = null role: RoleAssignmentRoleNestedFilter = null permission: PermissionNestedFilter = null username: StringFilter = null diff --git a/docs/manager/graphql-reference/v2-schema.graphql b/docs/manager/graphql-reference/v2-schema.graphql index a240aa4c3e5..12e1444493b 100644 --- a/docs/manager/graphql-reference/v2-schema.graphql +++ b/docs/manager/graphql-reference/v2-schema.graphql @@ -9809,7 +9809,7 @@ type RoleAssignmentEdge { """Added in 26.3.0. Filter for role assignments""" input RoleAssignmentFilter { - roleId: UUID = null + roleId: UUIDFilter = null role: RoleAssignmentRoleNestedFilter = null permission: PermissionNestedFilter = null username: StringFilter = null diff --git a/src/ai/backend/client/cli/v2/my/__init__.py b/src/ai/backend/client/cli/v2/my/__init__.py index 537c66bd4e9..2633b937160 100644 --- a/src/ai/backend/client/cli/v2/my/__init__.py +++ b/src/ai/backend/client/cli/v2/my/__init__.py @@ -21,6 +21,11 @@ def keypair() -> None: """My keypair commands.""" +@my.group(cls=LazyGroup, import_name="ai.backend.client.cli.v2.my.role:role") +def role() -> None: + """My role commands.""" + + @my.group( cls=LazyGroup, import_name="ai.backend.client.cli.v2.my.login_history:login_history", diff --git a/src/ai/backend/client/cli/v2/my/role.py b/src/ai/backend/client/cli/v2/my/role.py new file mode 100644 index 00000000000..d390756e04f --- /dev/null +++ b/src/ai/backend/client/cli/v2/my/role.py @@ -0,0 +1,85 @@ +"""CLI commands for self-service role operations.""" + +from __future__ import annotations + +import asyncio +from uuid import UUID + +import click + +from ai.backend.client.cli.v2.helpers import ( + create_v2_registry, + load_v2_config, + parse_order_options, + print_result, +) + + +@click.group() +def role() -> None: + """My role commands.""" + + +@role.command() +@click.option("--limit", default=None, type=int, help="Maximum number of results to return.") +@click.option("--offset", default=None, type=int, help="Number of results to skip.") +@click.option("--first", default=None, type=int, help="Cursor-based: return first N items.") +@click.option("--after", default=None, type=str, help="Cursor-based: return items after cursor.") +@click.option("--last", default=None, type=int, help="Cursor-based: return last N items.") +@click.option("--before", default=None, type=str, help="Cursor-based: return items before cursor.") +@click.option("--role-id", type=click.UUID, default=None, help="Filter by role UUID.") +@click.option( + "--order-by", + multiple=True, + help="Order by field:direction (e.g., granted_at:desc).", +) +def search( + limit: int | None, + offset: int | None, + first: int | None, + after: str | None, + last: int | None, + before: str | None, + role_id: UUID | None, + order_by: tuple[str, ...], +) -> None: + """Search my role assignments.""" + from ai.backend.common.dto.manager.query import UUIDFilter + from ai.backend.common.dto.manager.v2.rbac.request import ( + RoleAssignmentFilter, + RoleAssignmentOrderBy, + SearchRoleAssignmentsInput, + ) + from ai.backend.common.dto.manager.v2.rbac.types import RoleAssignmentOrderField + + filter_dto: RoleAssignmentFilter | None = None + if role_id is not None: + filter_dto = RoleAssignmentFilter( + role_id=UUIDFilter(equals=role_id), + ) + + orders = ( + parse_order_options(order_by, RoleAssignmentOrderField, RoleAssignmentOrderBy) + if order_by + else None + ) + + async def _run() -> None: + registry = await create_v2_registry(load_v2_config()) + try: + request = SearchRoleAssignmentsInput( + filter=filter_dto, + order=orders, + first=first, + after=after, + last=last, + before=before, + limit=limit, + offset=offset, + ) + result = await registry.rbac.my_search_assignments(request) + print_result(result) + finally: + await registry.close() + + asyncio.run(_run()) diff --git a/src/ai/backend/client/cli/v2/rbac/assignment.py b/src/ai/backend/client/cli/v2/rbac/assignment.py index 9acbb593a65..8b84ae60dbc 100644 --- a/src/ai/backend/client/cli/v2/rbac/assignment.py +++ b/src/ai/backend/client/cli/v2/rbac/assignment.py @@ -28,23 +28,23 @@ def assignment() -> None: multiple=True, help="Order by field:direction (e.g., username:asc, granted_at:desc).", ) -@click.option("--role-id", type=str, default=None, help="Filter by role UUID.") +@click.option("--role-id", type=click.UUID, default=None, help="Filter by role UUID.") @click.option("--username-contains", type=str, default=None, help="Filter by username (contains).") @click.option("--email-contains", type=str, default=None, help="Filter by email (contains).") def search( limit: int | None, offset: int | None, order_by: tuple[str, ...], - role_id: str | None, + role_id: UUID | None, username_contains: str | None, email_contains: str | None, ) -> None: """Search role assignments.""" - from ai.backend.common.dto.manager.query import StringFilter + from ai.backend.common.dto.manager.query import StringFilter, UUIDFilter from ai.backend.common.dto.manager.v2.rbac.request import ( - AdminSearchRoleAssignmentsGQLInput, RoleAssignmentFilter, RoleAssignmentOrderBy, + SearchRoleAssignmentsInput, ) from ai.backend.common.dto.manager.v2.rbac.types import RoleAssignmentOrderField @@ -52,7 +52,7 @@ def search( filter_dto: RoleAssignmentFilter | None = None if any([role_id is not None, username_contains is not None, email_contains is not None]): filter_dto = RoleAssignmentFilter( - role_id=UUID(role_id) if role_id is not None else None, + role_id=UUIDFilter(equals=role_id) if role_id is not None else None, username=( StringFilter(contains=username_contains) if username_contains is not None else None ), @@ -70,7 +70,7 @@ async def _run() -> None: registry = await create_v2_registry(load_v2_config()) try: result = await registry.rbac.search_assignments( - AdminSearchRoleAssignmentsGQLInput( + SearchRoleAssignmentsInput( filter=filter_dto, order=orders, limit=limit, diff --git a/src/ai/backend/client/v2/domains_v2/rbac.py b/src/ai/backend/client/v2/domains_v2/rbac.py index bc3ccbc89cb..963dc67b111 100644 --- a/src/ai/backend/client/v2/domains_v2/rbac.py +++ b/src/ai/backend/client/v2/domains_v2/rbac.py @@ -8,7 +8,6 @@ from ai.backend.common.dto.manager.v2.rbac.request import ( AdminSearchEntitiesGQLInput, AdminSearchPermissionsGQLInput, - AdminSearchRoleAssignmentsGQLInput, AssignRoleInput, BulkAssignRoleInput, BulkRevokeRoleInput, @@ -18,6 +17,7 @@ DeleteRoleInput, PurgeRoleInput, RevokeRoleInput, + SearchRoleAssignmentsInput, SearchRolesInput, UpdatePermissionInput, UpdateRoleInput, @@ -25,7 +25,6 @@ from ai.backend.common.dto.manager.v2.rbac.response import ( AdminSearchAssociationsPayload, AdminSearchPermissionsPayload, - AdminSearchRoleAssignmentsPayload, AdminSearchRolesPayload, BulkAssignRoleResultPayload, BulkRevokeRoleResultPayload, @@ -36,6 +35,7 @@ PurgeRolePayload, RoleAssignmentNode, RoleNode, + SearchRoleAssignmentsPayload, UpdateRolePayload, ) @@ -172,14 +172,25 @@ async def revoke_role(self, request: RevokeRoleInput) -> RoleAssignmentNode: ) async def search_assignments( - self, request: AdminSearchRoleAssignmentsGQLInput - ) -> AdminSearchRoleAssignmentsPayload: + self, request: SearchRoleAssignmentsInput + ) -> SearchRoleAssignmentsPayload: """Search role assignments with filters, orders, and pagination.""" return await self._client.typed_request( "POST", f"{_PATH}/assignments/search", request=request, - response_model=AdminSearchRoleAssignmentsPayload, + response_model=SearchRoleAssignmentsPayload, + ) + + async def my_search_assignments( + self, request: SearchRoleAssignmentsInput + ) -> SearchRoleAssignmentsPayload: + """Search role assignments for the current authenticated user.""" + return await self._client.typed_request( + "POST", + f"{_PATH}/assignments/my/search", + request=request, + response_model=SearchRoleAssignmentsPayload, ) async def bulk_assign_role(self, request: BulkAssignRoleInput) -> BulkAssignRoleResultPayload: diff --git a/src/ai/backend/common/dto/manager/v2/rbac/__init__.py b/src/ai/backend/common/dto/manager/v2/rbac/__init__.py index 9eb94f97604..97c871eae2d 100644 --- a/src/ai/backend/common/dto/manager/v2/rbac/__init__.py +++ b/src/ai/backend/common/dto/manager/v2/rbac/__init__.py @@ -5,7 +5,6 @@ from ai.backend.common.dto.manager.v2.rbac.request import ( AdminSearchEntitiesGQLInput, AdminSearchPermissionsGQLInput, - AdminSearchRoleAssignmentsGQLInput, AssignRoleInput, BulkAssignRoleInput, BulkRevokeRoleInput, @@ -24,6 +23,7 @@ RoleFilter, RoleNestedFilter, RoleOrderBy, + SearchRoleAssignmentsInput, SearchRolesInput, UpdatePermissionInput, UpdateRoleInput, @@ -93,7 +93,7 @@ # Input models (request) "AdminSearchEntitiesGQLInput", "AdminSearchPermissionsGQLInput", - "AdminSearchRoleAssignmentsGQLInput", + "SearchRoleAssignmentsInput", "SearchRolesInput", "AssignRoleInput", "BulkAssignRoleInput", diff --git a/src/ai/backend/common/dto/manager/v2/rbac/request.py b/src/ai/backend/common/dto/manager/v2/rbac/request.py index 6c32188520e..2cd2b3ac799 100644 --- a/src/ai/backend/common/dto/manager/v2/rbac/request.py +++ b/src/ai/backend/common/dto/manager/v2/rbac/request.py @@ -9,7 +9,7 @@ from pydantic import Field, field_validator from ai.backend.common.api_handlers import SENTINEL, BaseRequestModel, Sentinel -from ai.backend.common.dto.manager.query import DateTimeFilter, StringFilter +from ai.backend.common.dto.manager.query import DateTimeFilter, StringFilter, UUIDFilter from .types import ( OrderDirection, @@ -23,7 +23,7 @@ __all__ = ( "AdminSearchEntitiesGQLInput", "AdminSearchPermissionsGQLInput", - "AdminSearchRoleAssignmentsGQLInput", + "SearchRoleAssignmentsInput", "SearchRolesInput", "AssignRoleInput", "BulkAssignRoleInput", @@ -208,7 +208,7 @@ class PermissionNestedFilter(BaseRequestModel): class RoleAssignmentFilter(BaseRequestModel): """Filter for role assignments.""" - role_id: UUID | None = None + role_id: UUIDFilter | None = None role: RoleNestedFilter | None = None permission: PermissionNestedFilter | None = None username: StringFilter | None = None @@ -303,8 +303,8 @@ class SearchRolesInput(BaseRequestModel): offset: int | None = None -class AdminSearchRoleAssignmentsGQLInput(BaseRequestModel): - """GQL pagination search input for role assignments.""" +class SearchRoleAssignmentsInput(BaseRequestModel): + """Pagination search input for role assignments.""" filter: RoleAssignmentFilter | None = None order: list[RoleAssignmentOrderBy] | None = None diff --git a/src/ai/backend/common/dto/manager/v2/rbac/response.py b/src/ai/backend/common/dto/manager/v2/rbac/response.py index 02c7569c640..2fa548dbf89 100644 --- a/src/ai/backend/common/dto/manager/v2/rbac/response.py +++ b/src/ai/backend/common/dto/manager/v2/rbac/response.py @@ -21,7 +21,7 @@ __all__ = ( "AdminSearchAssociationsPayload", "AdminSearchPermissionsPayload", - "AdminSearchRoleAssignmentsPayload", + "SearchRoleAssignmentsPayload", "AdminSearchRolesPayload", "AssociationScopesEntitiesNode", "BulkAssignRoleFailureInfo", @@ -183,7 +183,7 @@ class AdminSearchPermissionsPayload(BaseResponseModel): has_previous_page: bool = Field(description="Whether there is a previous page.") -class AdminSearchRoleAssignmentsPayload(BaseResponseModel): +class SearchRoleAssignmentsPayload(BaseResponseModel): """Paginated result for role assignment search.""" items: list[RoleAssignmentNode] = Field(description="List of role assignment nodes.") diff --git a/src/ai/backend/manager/api/adapters/rbac.py b/src/ai/backend/manager/api/adapters/rbac.py index 966259d9519..36fef958c38 100644 --- a/src/ai/backend/manager/api/adapters/rbac.py +++ b/src/ai/backend/manager/api/adapters/rbac.py @@ -10,6 +10,7 @@ from uuid import UUID from ai.backend.common.api_handlers import SENTINEL +from ai.backend.common.contexts.user import current_user from ai.backend.common.data.filter_specs import StringMatchSpec from ai.backend.common.data.permission.types import OperationType as InternalOperationType from ai.backend.common.data.permission.types import RBACElementType @@ -51,7 +52,7 @@ from ai.backend.common.dto.manager.v2.rbac.request import ( AdminSearchEntitiesGQLInput, AdminSearchPermissionsGQLInput, - AdminSearchRoleAssignmentsGQLInput, + SearchRoleAssignmentsInput, SearchRolesInput, ) from ai.backend.common.dto.manager.v2.rbac.request import ( @@ -111,6 +112,7 @@ from ai.backend.common.dto.manager.v2.rbac.types import ( OrderDirection as OrderDirectionV2, ) +from ai.backend.common.exception import UnreachableError from ai.backend.manager.actions.action import build_operation_description from ai.backend.manager.api.adapters.pagination import PaginationSpec from ai.backend.manager.data.common.types import SearchResult @@ -680,12 +682,33 @@ async def search_roles_in_scope( has_previous_page=raw.has_previous_page, ) - async def admin_search_role_assignments_gql( + async def my_search_role_assignments( self, - input: AdminSearchRoleAssignmentsGQLInput, + input: SearchRoleAssignmentsInput, + ) -> SearchResult[RoleAssignmentNode]: + """Search role assignments for the current authenticated user.""" + me = current_user() + if me is None: + raise UnreachableError("User context is not available") + return await self._search_role_assignments( + input, + base_conditions=[AssignedUserConditions.by_user_id(me.user_id)], + ) + + async def admin_search_role_assignments( + self, + input: SearchRoleAssignmentsInput, base_conditions: Sequence[QueryCondition] | None = None, ) -> SearchResult[RoleAssignmentNode]: - """Search role assignments with cursor/offset pagination.""" + """Search role assignments with cursor/offset pagination (admin).""" + return await self._search_role_assignments(input, base_conditions=base_conditions) + + async def _search_role_assignments( + self, + input: SearchRoleAssignmentsInput, + base_conditions: Sequence[QueryCondition] | None = None, + ) -> SearchResult[RoleAssignmentNode]: + """Internal implementation for searching role assignments.""" conditions = self._convert_assignment_filter(input.filter) if input.filter else [] orders = self._convert_assignment_orders(input.order) if input.order else [] querier = self._build_querier( @@ -1242,7 +1265,13 @@ def _convert_permission_nested_filter( def _convert_assignment_filter(self, f: RoleAssignmentFilterDTO) -> list[QueryCondition]: conditions: list[QueryCondition] = [] if f.role_id is not None: - conditions.append(AssignedUserConditions.by_role_id(f.role_id)) + condition = self.convert_uuid_filter( + f.role_id, + equals_factory=AssignedUserConditions.by_role_id_equals, + in_factory=AssignedUserConditions.by_role_id_in, + ) + if condition is not None: + conditions.append(condition) if f.role is not None: conditions.extend(self._convert_role_nested_filter(f.role)) if f.permission is not None: diff --git a/src/ai/backend/manager/api/gql/rbac/resolver/role.py b/src/ai/backend/manager/api/gql/rbac/resolver/role.py index 924edccb294..a92120c2b12 100644 --- a/src/ai/backend/manager/api/gql/rbac/resolver/role.py +++ b/src/ai/backend/manager/api/gql/rbac/resolver/role.py @@ -10,7 +10,7 @@ from ai.backend.common.contexts.user import current_user from ai.backend.common.data.permission.types import RBACElementType from ai.backend.common.dto.manager.v2.rbac.request import ( - AdminSearchRoleAssignmentsGQLInput, + SearchRoleAssignmentsInput, SearchRolesInput, ) from ai.backend.manager.api.gql.base import encode_cursor @@ -125,8 +125,8 @@ async def admin_role_assignments( offset: int | None = None, ) -> RoleAssignmentConnection: check_admin_only() - result = await info.context.adapters.rbac.admin_search_role_assignments_gql( - AdminSearchRoleAssignmentsGQLInput( + result = await info.context.adapters.rbac.admin_search_role_assignments( + SearchRoleAssignmentsInput( filter=filter.to_pydantic() if filter is not None else None, order=[o.to_pydantic() for o in order_by] if order_by is not None else None, first=first, @@ -178,8 +178,8 @@ async def my_roles( raise InsufficientPrivilege("Authentication required") - result = await info.context.adapters.rbac.admin_search_role_assignments_gql( - AdminSearchRoleAssignmentsGQLInput( + result = await info.context.adapters.rbac.admin_search_role_assignments( + SearchRoleAssignmentsInput( filter=filter.to_pydantic() if filter is not None else None, order=[o.to_pydantic() for o in order_by] if order_by is not None else None, first=first, diff --git a/src/ai/backend/manager/api/gql/rbac/types/role.py b/src/ai/backend/manager/api/gql/rbac/types/role.py index f423fc3fb6b..cddd7bc9db9 100644 --- a/src/ai/backend/manager/api/gql/rbac/types/role.py +++ b/src/ai/backend/manager/api/gql/rbac/types/role.py @@ -16,7 +16,7 @@ from ai.backend.common.dto.manager.v2.rbac.request import ( AdminSearchEntitiesGQLInput, AdminSearchPermissionsGQLInput, - AdminSearchRoleAssignmentsGQLInput, + SearchRoleAssignmentsInput, ) from ai.backend.common.dto.manager.v2.rbac.request import ( AssignRoleInput as AssignRoleInputDTO, @@ -88,7 +88,7 @@ from ai.backend.common.dto.manager.v2.rbac.types import ( RoleStatusFilter as RoleStatusFilterDTO, ) -from ai.backend.manager.api.gql.base import OrderDirection, StringFilter, encode_cursor +from ai.backend.manager.api.gql.base import OrderDirection, StringFilter, UUIDFilter, encode_cursor from ai.backend.manager.api.gql.decorators import ( BackendAIGQLMeta, PydanticInputMixin, @@ -266,7 +266,7 @@ async def users( offset: int | None = None, ) -> RoleAssignmentConnection: # Add role_id filter to scope assignments to this role - role_filter = RoleAssignmentFilter(role_id=UUID(self.id)) + role_filter = RoleAssignmentFilter(role_id=UUIDFilter(equals=UUID(self.id))) if filter is not None: # Merge with user-provided filter combined_filter = RoleAssignmentFilter( @@ -281,8 +281,8 @@ async def users( pydantic_filter = combined_filter.to_pydantic() if combined_filter is not None else None pydantic_order = [o.to_pydantic() for o in order_by] if order_by is not None else None - result = await info.context.adapters.rbac.admin_search_role_assignments_gql( - AdminSearchRoleAssignmentsGQLInput( + result = await info.context.adapters.rbac.admin_search_role_assignments( + SearchRoleAssignmentsInput( filter=pydantic_filter, order=pydantic_order, first=first, @@ -532,7 +532,7 @@ class RoleAssignmentRoleNestedFilterGQL(PydanticInputMixin[RoleNestedFilterDTO]) name="RoleAssignmentFilter", ) class RoleAssignmentFilter(PydanticInputMixin[RoleAssignmentFilterDTO], GQLFilter): - role_id: UUID | None = None + role_id: UUIDFilter | None = None role: RoleAssignmentRoleNestedFilterGQL | None = None permission: ( Annotated[ diff --git a/src/ai/backend/manager/api/rest/v2/rbac/handler.py b/src/ai/backend/manager/api/rest/v2/rbac/handler.py index 85ac9981215..499c09bbf00 100644 --- a/src/ai/backend/manager/api/rest/v2/rbac/handler.py +++ b/src/ai/backend/manager/api/rest/v2/rbac/handler.py @@ -11,7 +11,6 @@ from ai.backend.common.dto.manager.v2.rbac.request import ( AdminSearchEntitiesGQLInput, AdminSearchPermissionsGQLInput, - AdminSearchRoleAssignmentsGQLInput, AssignRoleInput, BulkAssignRoleInput, BulkRevokeRoleInput, @@ -21,6 +20,7 @@ DeleteRoleInput, PurgeRoleInput, RevokeRoleInput, + SearchRoleAssignmentsInput, SearchRolesInput, UpdatePermissionInput, UpdateRoleInput, @@ -28,9 +28,9 @@ from ai.backend.common.dto.manager.v2.rbac.response import ( AdminSearchAssociationsPayload, AdminSearchPermissionsPayload, - AdminSearchRoleAssignmentsPayload, AdminSearchRolesPayload, ScopeEntityOperationCombinationInfo, + SearchRoleAssignmentsPayload, ) from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.api.rest.v2.path_params import ProjectIdPathParam, RoleIdPathParam @@ -191,11 +191,25 @@ async def revoke_role( async def search_assignments( self, - body: BodyParam[AdminSearchRoleAssignmentsGQLInput], + body: BodyParam[SearchRoleAssignmentsInput], ) -> APIResponse: """Search role assignments with filters, orders, and pagination.""" - result = await self._adapter.admin_search_role_assignments_gql(body.parsed) - payload = AdminSearchRoleAssignmentsPayload( + result = await self._adapter.admin_search_role_assignments(body.parsed) + payload = SearchRoleAssignmentsPayload( + items=result.items, + total_count=result.total_count, + has_next_page=result.has_next_page, + has_previous_page=result.has_previous_page, + ) + return APIResponse.build(status_code=HTTPStatus.OK, response_model=payload) + + async def my_search_assignments( + self, + body: BodyParam[SearchRoleAssignmentsInput], + ) -> APIResponse: + """Search role assignments for the current authenticated user.""" + result = await self._adapter.my_search_role_assignments(body.parsed) + payload = SearchRoleAssignmentsPayload( items=result.items, total_count=result.total_count, has_next_page=result.has_next_page, diff --git a/src/ai/backend/manager/api/rest/v2/rbac/registry.py b/src/ai/backend/manager/api/rest/v2/rbac/registry.py index 73fc1a7eb24..b290b860d5c 100644 --- a/src/ai/backend/manager/api/rest/v2/rbac/registry.py +++ b/src/ai/backend/manager/api/rest/v2/rbac/registry.py @@ -109,6 +109,12 @@ def register_v2_rbac_routes( handler.search_assignments, middlewares=[superadmin_required], ) + registry.add( + "POST", + "/assignments/my/search", + handler.my_search_assignments, + middlewares=[auth_required], + ) registry.add( "POST", "/assignments/bulk-assign", diff --git a/src/ai/backend/manager/models/rbac_models/conditions.py b/src/ai/backend/manager/models/rbac_models/conditions.py index 1464ed341a9..11d2622fa93 100644 --- a/src/ai/backend/manager/models/rbac_models/conditions.py +++ b/src/ai/backend/manager/models/rbac_models/conditions.py @@ -8,7 +8,12 @@ import sqlalchemy as sa -from ai.backend.common.data.filter_specs import StringInMatchSpec, StringMatchSpec +from ai.backend.common.data.filter_specs import ( + StringInMatchSpec, + StringMatchSpec, + UUIDEqualMatchSpec, + UUIDInMatchSpec, +) from ai.backend.common.data.permission.types import RBACElementType from ai.backend.manager.data.permission.id import ObjectId from ai.backend.manager.data.permission.status import RoleStatus @@ -255,6 +260,26 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner + @staticmethod + def by_role_id_equals(spec: UUIDEqualMatchSpec) -> QueryCondition: + def inner() -> sa.sql.expression.ColumnElement[bool]: + condition = UserRoleRow.role_id == spec.value + if spec.negated: + condition = sa.not_(condition) + return condition + + return inner + + @staticmethod + def by_role_id_in(spec: UUIDInMatchSpec) -> QueryCondition: + def inner() -> sa.sql.expression.ColumnElement[bool]: + condition = UserRoleRow.role_id.in_(spec.values) + if spec.negated: + condition = sa.not_(condition) + return condition + + return inner + @staticmethod def by_username_contains(spec: StringMatchSpec) -> QueryCondition: def inner() -> sa.sql.expression.ColumnElement[bool]: diff --git a/tests/component/rbac/test_rbac_role_v2.py b/tests/component/rbac/test_rbac_role_v2.py new file mode 100644 index 00000000000..a2cd46528de --- /dev/null +++ b/tests/component/rbac/test_rbac_role_v2.py @@ -0,0 +1,289 @@ +"""Component tests for RBAC role assignment via v2 REST API.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, MagicMock + +import pytest +import yarl + +from ai.backend.client.v2.auth import HMACAuth +from ai.backend.client.v2.config import ClientConfig +from ai.backend.client.v2.exceptions import PermissionDeniedError +from ai.backend.client.v2.v2_registry import V2ClientRegistry +from ai.backend.common.dto.manager.v2.rbac.request import ( + AssignRoleInput, + RevokeRoleInput, + SearchRoleAssignmentsInput, +) +from ai.backend.common.dto.manager.v2.rbac.response import ( + RoleAssignmentNode, + SearchRoleAssignmentsPayload, +) +from ai.backend.manager.api.adapters.rbac import RBACAdapter +from ai.backend.manager.api.rest.admin.handler import AdminHandler +from ai.backend.manager.api.rest.admin.registry import register_admin_routes +from ai.backend.manager.api.rest.rbac.handler import RBACHandler +from ai.backend.manager.api.rest.rbac.registry import register_rbac_routes +from ai.backend.manager.api.rest.routing import RouteRegistry +from ai.backend.manager.api.rest.types import RouteDeps +from ai.backend.manager.api.rest.v2.rbac.handler import V2RBACHandler +from ai.backend.manager.api.rest.v2.rbac.registry import register_v2_rbac_routes +from ai.backend.manager.models.utils import ExtendedAsyncSAEngine +from ai.backend.manager.repositories.permission_controller.repository import ( + PermissionControllerRepository, +) +from ai.backend.manager.services.permission_contoller.processors import ( + PermissionControllerProcessors, +) +from ai.backend.manager.services.permission_contoller.service import PermissionControllerService + +if TYPE_CHECKING: + from tests.component.conftest import ServerInfo, UserFixtureData + + from ai.backend.client.v2.registry import BackendAIClientRegistry + from ai.backend.common.dto.manager.rbac.response import CreateRoleResponse + + +@pytest.fixture() +def permission_controller_processors( + database_engine: ExtendedAsyncSAEngine, +) -> PermissionControllerProcessors: + repo = PermissionControllerRepository(database_engine) + service = PermissionControllerService( + repo, group_repository=MagicMock(), rbac_action_registry=[] + ) + validators = MagicMock() + validators.rbac.scope.validate = AsyncMock() + return PermissionControllerProcessors( + service=service, action_monitors=[], validators=validators + ) + + +@pytest.fixture() +def server_module_registries( + route_deps: RouteDeps, + permission_controller_processors: PermissionControllerProcessors, +) -> list[RouteRegistry]: + """Register both v1 RBAC (for role creation) and v2 RBAC (for assignments) routes.""" + rbac_registry = register_rbac_routes( + RBACHandler(permission_controller=permission_controller_processors), route_deps + ) + admin_registry = register_admin_routes( + AdminHandler(gql_schema=MagicMock(), gql_deps=MagicMock(), strawberry_schema=MagicMock()), + route_deps, + sub_registries=[rbac_registry], + gql_ws_handler=MagicMock(), + ) + + processors = MagicMock() + processors.permission_controller = permission_controller_processors + adapter = RBACAdapter(processors) + handler = V2RBACHandler(adapter=adapter) + v2_reg = RouteRegistry.create("v2", route_deps.cors_options) + v2_reg.add_subregistry(register_v2_rbac_routes(handler, route_deps)) + + return [admin_registry, v2_reg] + + +@pytest.fixture() +async def admin_v2_registry( + server: ServerInfo, + admin_user_fixture: UserFixtureData, +) -> AsyncIterator[V2ClientRegistry]: + """Create a V2ClientRegistry with superadmin keypair for v2 REST endpoints.""" + registry = await V2ClientRegistry.create( + ClientConfig(endpoint=yarl.URL(server.url)), + HMACAuth( + access_key=admin_user_fixture.keypair.access_key, + secret_key=admin_user_fixture.keypair.secret_key, + ), + ) + try: + yield registry + finally: + await registry.close() + + +@pytest.fixture() +async def user_v2_registry( + server: ServerInfo, + regular_user_fixture: UserFixtureData, +) -> AsyncIterator[V2ClientRegistry]: + """Create a V2ClientRegistry with regular-user keypair for v2 REST endpoints.""" + registry = await V2ClientRegistry.create( + ClientConfig(endpoint=yarl.URL(server.url)), + HMACAuth( + access_key=regular_user_fixture.keypair.access_key, + secret_key=regular_user_fixture.keypair.secret_key, + ), + ) + try: + yield registry + finally: + await registry.close() + + +class TestAssignRoleV2: + """Test role assignment via v2 REST API.""" + + async def test_admin_assigns_role( + self, + admin_v2_registry: V2ClientRegistry, + admin_registry: BackendAIClientRegistry, + target_role: CreateRoleResponse, + admin_user_fixture: UserFixtureData, + ) -> None: + result = await admin_v2_registry.rbac.assign_role( + AssignRoleInput( + user_id=admin_user_fixture.user_uuid, + role_id=target_role.role.id, + ) + ) + assert isinstance(result, RoleAssignmentNode) + assert result.user_id == admin_user_fixture.user_uuid + assert result.role_id == target_role.role.id + + # Clean up + await admin_v2_registry.rbac.revoke_role( + RevokeRoleInput( + user_id=admin_user_fixture.user_uuid, + role_id=target_role.role.id, + ) + ) + + async def test_regular_user_cannot_assign_role( + self, + user_v2_registry: V2ClientRegistry, + admin_registry: BackendAIClientRegistry, + target_role: CreateRoleResponse, + regular_user_fixture: UserFixtureData, + ) -> None: + with pytest.raises(PermissionDeniedError): + await user_v2_registry.rbac.assign_role( + AssignRoleInput( + user_id=regular_user_fixture.user_uuid, + role_id=target_role.role.id, + ) + ) + + +class TestRevokeRoleV2: + """Test role revocation via v2 REST API.""" + + async def test_admin_revokes_role( + self, + admin_v2_registry: V2ClientRegistry, + admin_registry: BackendAIClientRegistry, + target_role: CreateRoleResponse, + admin_user_fixture: UserFixtureData, + ) -> None: + # Assign first + await admin_v2_registry.rbac.assign_role( + AssignRoleInput( + user_id=admin_user_fixture.user_uuid, + role_id=target_role.role.id, + ) + ) + # Revoke + result = await admin_v2_registry.rbac.revoke_role( + RevokeRoleInput( + user_id=admin_user_fixture.user_uuid, + role_id=target_role.role.id, + ) + ) + assert isinstance(result, RoleAssignmentNode) + assert result.user_id == admin_user_fixture.user_uuid + assert result.role_id == target_role.role.id + + async def test_regular_user_cannot_revoke_role( + self, + user_v2_registry: V2ClientRegistry, + admin_registry: BackendAIClientRegistry, + target_role: CreateRoleResponse, + regular_user_fixture: UserFixtureData, + ) -> None: + with pytest.raises(PermissionDeniedError): + await user_v2_registry.rbac.revoke_role( + RevokeRoleInput( + user_id=regular_user_fixture.user_uuid, + role_id=target_role.role.id, + ) + ) + + +class TestMySearchAssignmentsV2: + """Test self-service role assignment search via v2 REST API.""" + + async def test_user_searches_own_assignments( + self, + admin_v2_registry: V2ClientRegistry, + user_v2_registry: V2ClientRegistry, + admin_registry: BackendAIClientRegistry, + target_role: CreateRoleResponse, + regular_user_fixture: UserFixtureData, + ) -> None: + # Admin assigns a role to the regular user + await admin_v2_registry.rbac.assign_role( + AssignRoleInput( + user_id=regular_user_fixture.user_uuid, + role_id=target_role.role.id, + ) + ) + try: + # Regular user searches their own assignments + result = await user_v2_registry.rbac.my_search_assignments(SearchRoleAssignmentsInput()) + assert isinstance(result, SearchRoleAssignmentsPayload) + assert result.total_count >= 1 + assert any(a.role_id == target_role.role.id for a in result.items) + # All results should belong to the current user + assert all(a.user_id == regular_user_fixture.user_uuid for a in result.items) + finally: + # Clean up + await admin_v2_registry.rbac.revoke_role( + RevokeRoleInput( + user_id=regular_user_fixture.user_uuid, + role_id=target_role.role.id, + ) + ) + + async def test_user_does_not_see_other_users_assignments( + self, + admin_v2_registry: V2ClientRegistry, + user_v2_registry: V2ClientRegistry, + admin_registry: BackendAIClientRegistry, + target_role: CreateRoleResponse, + admin_user_fixture: UserFixtureData, + regular_user_fixture: UserFixtureData, + ) -> None: + # Admin assigns a role to the admin user (not the regular user) + await admin_v2_registry.rbac.assign_role( + AssignRoleInput( + user_id=admin_user_fixture.user_uuid, + role_id=target_role.role.id, + ) + ) + try: + # Regular user searches — should NOT see admin's assignment + result = await user_v2_registry.rbac.my_search_assignments(SearchRoleAssignmentsInput()) + assert all(a.user_id != admin_user_fixture.user_uuid for a in result.items) + finally: + # Clean up + await admin_v2_registry.rbac.revoke_role( + RevokeRoleInput( + user_id=admin_user_fixture.user_uuid, + role_id=target_role.role.id, + ) + ) + + async def test_empty_result_when_no_roles_assigned( + self, + user_v2_registry: V2ClientRegistry, + admin_registry: BackendAIClientRegistry, + ) -> None: + result = await user_v2_registry.rbac.my_search_assignments(SearchRoleAssignmentsInput()) + assert isinstance(result, SearchRoleAssignmentsPayload) + assert result.total_count == 0 + assert result.items == [] diff --git a/tests/unit/manager/api/gql/rbac/test_my_roles_resolver.py b/tests/unit/manager/api/gql/rbac/test_my_roles_resolver.py index 1495d0b98f9..0f89f835ca4 100644 --- a/tests/unit/manager/api/gql/rbac/test_my_roles_resolver.py +++ b/tests/unit/manager/api/gql/rbac/test_my_roles_resolver.py @@ -42,7 +42,7 @@ async def test_calls_adapter_with_user_condition( ) ) info = MagicMock() - info.context.adapters.rbac.admin_search_role_assignments_gql = mock_search + info.context.adapters.rbac.admin_search_role_assignments = mock_search with patch( "ai.backend.manager.api.gql.rbac.resolver.role.current_user", @@ -105,7 +105,7 @@ async def test_passes_pagination_params( ) ) info = MagicMock() - info.context.adapters.rbac.admin_search_role_assignments_gql = mock_search + info.context.adapters.rbac.admin_search_role_assignments = mock_search with patch( "ai.backend.manager.api.gql.rbac.resolver.role.current_user",