diff --git a/changes/10320.feature.md b/changes/10320.feature.md new file mode 100644 index 00000000000..6155ea4aca5 --- /dev/null +++ b/changes/10320.feature.md @@ -0,0 +1 @@ +Support resolving session entities in RBAC entity and permission scope queries diff --git a/docs/manager/graphql-reference/supergraph.graphql b/docs/manager/graphql-reference/supergraph.graphql index 88989077722..3ba4f7266df 100644 --- a/docs/manager/graphql-reference/supergraph.graphql +++ b/docs/manager/graphql-reference/supergraph.graphql @@ -4639,6 +4639,7 @@ union EntityNode @join__unionMember(graph: STRAWBERRY, member: "VirtualFolderNode") @join__unionMember(graph: STRAWBERRY, member: "ImageV2") @join__unionMember(graph: STRAWBERRY, member: "ComputeSessionNode") + @join__unionMember(graph: STRAWBERRY, member: "SessionV2") @join__unionMember(graph: STRAWBERRY, member: "Artifact") @join__unionMember(graph: STRAWBERRY, member: "ArtifactRegistry") @join__unionMember(graph: STRAWBERRY, member: "AppConfig") @@ -4649,7 +4650,7 @@ union EntityNode @join__unionMember(graph: STRAWBERRY, member: "ContainerRegistryV2") @join__unionMember(graph: STRAWBERRY, member: "ArtifactRevision") @join__unionMember(graph: STRAWBERRY, member: "Role") - = UserV2 | ProjectV2 | DomainV2 | VirtualFolderNode | ImageV2 | ComputeSessionNode | Artifact | ArtifactRegistry | AppConfig | NotificationChannel | NotificationRule | ModelDeployment | ResourceGroup | ContainerRegistryV2 | ArtifactRevision | Role + = UserV2 | ProjectV2 | DomainV2 | VirtualFolderNode | ImageV2 | ComputeSessionNode | SessionV2 | Artifact | ArtifactRegistry | AppConfig | NotificationChannel | NotificationRule | ModelDeployment | ResourceGroup | ContainerRegistryV2 | ArtifactRevision | Role """Added in 26.3.0. Order by specification for entity associations""" input EntityOrderBy diff --git a/docs/manager/graphql-reference/v2-schema.graphql b/docs/manager/graphql-reference/v2-schema.graphql index f334b31750e..75f281f342a 100644 --- a/docs/manager/graphql-reference/v2-schema.graphql +++ b/docs/manager/graphql-reference/v2-schema.graphql @@ -2677,7 +2677,7 @@ input EntityFilter { NOT: [EntityFilter!] = null } -union EntityNode = UserV2 | ProjectV2 | DomainV2 | VirtualFolderNode | ImageV2 | ComputeSessionNode | Artifact | ArtifactRegistry | AppConfig | NotificationChannel | NotificationRule | ModelDeployment | ResourceGroup | ContainerRegistryV2 | ArtifactRevision | Role +union EntityNode = UserV2 | ProjectV2 | DomainV2 | VirtualFolderNode | ImageV2 | ComputeSessionNode | SessionV2 | Artifact | ArtifactRegistry | AppConfig | NotificationChannel | NotificationRule | ModelDeployment | ResourceGroup | ContainerRegistryV2 | ArtifactRevision | Role """Added in 26.3.0. Order by specification for entity associations""" input EntityOrderBy { diff --git a/src/ai/backend/manager/api/gql/rbac/types/entity.py b/src/ai/backend/manager/api/gql/rbac/types/entity.py index f58db6ba626..b8395eab527 100644 --- a/src/ai/backend/manager/api/gql/rbac/types/entity.py +++ b/src/ai/backend/manager/api/gql/rbac/types/entity.py @@ -63,7 +63,7 @@ async def entity( *, info: Info[StrawberryGQLContext], ) -> EntityNode | None: - from ai.backend.common.types import ImageID + from ai.backend.common.types import ImageID, SessionId from ai.backend.manager.api.gql.artifact.types import ArtifactRevision from ai.backend.manager.api.gql.container_registry.types import ContainerRegistryGQL from ai.backend.manager.api.gql.deployment.types.deployment import ModelDeployment @@ -76,6 +76,7 @@ async def entity( from ai.backend.manager.api.gql.project_v2.types.node import ProjectV2GQL from ai.backend.manager.api.gql.rbac.types.role import RoleGQL from ai.backend.manager.api.gql.resource_group.types import ResourceGroupGQL + from ai.backend.manager.api.gql.session.types import SessionV2GQL from ai.backend.manager.api.gql.user.types.node import UserV2GQL element_type = self.entity_type.to_element() @@ -146,9 +147,15 @@ async def entity( if cr_data is None: return None return ContainerRegistryGQL.from_data(cr_data) + case RBACElementType.SESSION: + session_data = await data_loaders.session_loader.load( + SessionId(uuid.UUID(self.entity_id)) + ) + if session_data is None: + return None + return SessionV2GQL.from_data(session_data) case ( - RBACElementType.SESSION - | RBACElementType.VFOLDER + RBACElementType.VFOLDER | RBACElementType.KEYPAIR | RBACElementType.NETWORK | RBACElementType.STORAGE_HOST diff --git a/src/ai/backend/manager/api/gql/rbac/types/entity_node.py b/src/ai/backend/manager/api/gql/rbac/types/entity_node.py index 95e8936c785..11682471513 100644 --- a/src/ai/backend/manager/api/gql/rbac/types/entity_node.py +++ b/src/ai/backend/manager/api/gql/rbac/types/entity_node.py @@ -24,6 +24,7 @@ from ai.backend.manager.api.gql.project_v2.types.node import ProjectV2GQL from ai.backend.manager.api.gql.rbac.types.role import RoleGQL from ai.backend.manager.api.gql.resource_group.types import ResourceGroupGQL +from ai.backend.manager.api.gql.session.types import SessionV2GQL from ai.backend.manager.api.gql.session_federation import Session from ai.backend.manager.api.gql.user.types.node import UserV2GQL from ai.backend.manager.api.gql.vfolder import VFolder @@ -39,6 +40,7 @@ | VFolder | ImageV2GQL | Session + | SessionV2GQL | Artifact | ArtifactRegistry | AppConfig diff --git a/src/ai/backend/manager/api/gql/rbac/types/permission.py b/src/ai/backend/manager/api/gql/rbac/types/permission.py index b5abd34f973..526593919de 100644 --- a/src/ai/backend/manager/api/gql/rbac/types/permission.py +++ b/src/ai/backend/manager/api/gql/rbac/types/permission.py @@ -15,6 +15,7 @@ OperationType, RBACElementType, ) +from ai.backend.common.types import SessionId from ai.backend.manager.api.gql.base import OrderDirection from ai.backend.manager.api.gql.rbac.types.entity_node import EntityNode from ai.backend.manager.api.gql.types import GQLFilter, GQLOrderBy, StrawberryGQLContext @@ -194,6 +195,7 @@ async def scope( from ai.backend.manager.api.gql.project_v2.types.node import ProjectV2GQL from ai.backend.manager.api.gql.rbac.types.role import RoleGQL from ai.backend.manager.api.gql.resource_group.types import ResourceGroupGQL + from ai.backend.manager.api.gql.session.types import SessionV2GQL from ai.backend.manager.api.gql.user.types.node import UserV2GQL element_type = self.scope_type.to_element() @@ -243,9 +245,15 @@ async def scope( if cr_data is None: return None return ContainerRegistryGQL.from_data(cr_data) + case RBACElementType.SESSION: + session_data = await data_loaders.session_loader.load( + SessionId(uuid.UUID(self.scope_id)) + ) + if session_data is None: + return None + return SessionV2GQL.from_data(session_data) case ( - RBACElementType.SESSION - | RBACElementType.VFOLDER + RBACElementType.VFOLDER | RBACElementType.KEYPAIR | RBACElementType.NOTIFICATION_CHANNEL | RBACElementType.NETWORK