Skip to content

Commit

Permalink
feat: Implement per-project images API based on RBAC
Browse files Browse the repository at this point in the history
  • Loading branch information
jopemachine committed Nov 12, 2024
1 parent 8289520 commit b14387f
Show file tree
Hide file tree
Showing 3 changed files with 327 additions and 21 deletions.
54 changes: 54 additions & 0 deletions src/ai/backend/manager/api/schema.graphql
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,29 @@ type Queries {

"""Added in 24.03.1"""
customized_images: [ImageNode]

"""Added in 24.12.0."""
image_node(
id: GlobalIDField!

"""Default is read_attribute."""
permission: ImagePermissionValueField = "read_attribute"
): ImageNode

"""Added in 24.12.0."""
image_nodes(
scope_id: ScopeField!

"""Default is read_attribute."""
permission: ImagePermissionValueField = "read_attribute"
filter: String
order: String
offset: Int
before: String
after: String
first: Int
last: Int
): ImageConnection
user(domain_name: String, email: String): User
user_from_uuid(domain_name: String, user_id: ID): User
users(domain_name: String, group_id: UUID, is_active: Boolean, status: String): [User]
Expand Down Expand Up @@ -337,6 +360,11 @@ type ImageNode implements Node {

"""Added in 24.03.4. The array of image aliases."""
aliases: [String]

"""
Added in 24.12.0. One of ['read_attribute', 'update_attribute', 'create_container', 'forget_image'].
"""
permissions: [ImagePermissionValueField]
}

type KVPair {
Expand All @@ -356,6 +384,11 @@ type ResourceLimit {
max: String
}

"""
Added in 24.12.0. One of ['read_attribute', 'update_attribute', 'create_container', 'forget_image'].
"""
scalar ImagePermissionValueField

type AgentList implements PaginatedList {
items: [Agent]!
total_count: Int!
Expand Down Expand Up @@ -765,6 +798,27 @@ type Image {
hash: String
}

"""Added in 24.12.0."""
type ImageConnection {
"""Pagination data for this connection."""
pageInfo: PageInfo!

"""Contains the nodes in this connection."""
edges: [ImageEdge]!

"""Total count of the GQL nodes of the query."""
count: Int
}

"""Added in 24.12.0. A Relay edge containing a `Image` and its cursor."""
type ImageEdge {
"""The item at the end of the edge"""
node: ImageNode

"""A cursor for use in pagination"""
cursor: String!
}

type User implements Item {
id: ID
uuid: UUID
Expand Down
66 changes: 65 additions & 1 deletion src/ai/backend/manager/models/gql.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,9 @@
ForgetImage,
ForgetImageById,
Image,
ImageConnection,
ImageNode,
ImagePermissionValueField,
ModifyImage,
PreloadImage,
RescanImages,
Expand Down Expand Up @@ -135,7 +137,12 @@
)
from .keypair import CreateKeyPair, DeleteKeyPair, KeyPair, KeyPairList, ModifyKeyPair
from .rbac import ScopeType, SystemScope
from .rbac.permission_defs import AgentPermission, ComputeSessionPermission, DomainPermission
from .rbac.permission_defs import (
AgentPermission,
ComputeSessionPermission,
DomainPermission,
ImagePermission,
)
from .rbac.permission_defs import VFolderPermission as VFolderRBACPermission
from .resource_policy import (
CreateKeyPairResourcePolicy,
Expand Down Expand Up @@ -517,6 +524,25 @@ class Queries(graphene.ObjectType):

customized_images = graphene.List(ImageNode, description="Added in 24.03.1")

image_node = graphene.Field(
ImageNode,
description="Added in 24.12.0.",
id=GlobalIDField(required=True),
permission=ImagePermissionValueField(
default_value=ImagePermission.READ_ATTRIBUTE,
description=f"Default is {ImagePermission.READ_ATTRIBUTE.value}.",
),
)
image_nodes = PaginatedConnectionField(
ImageConnection,
description="Added in 24.12.0.",
scope_id=ScopeField(required=True),
permission=ImagePermissionValueField(
default_value=ImagePermission.READ_ATTRIBUTE,
description=f"Default is {ImagePermission.READ_ATTRIBUTE.value}.",
),
)

user = graphene.Field(
User,
domain_name=graphene.String(),
Expand Down Expand Up @@ -1570,6 +1596,44 @@ async def resolve_user_nodes(
last,
)

@staticmethod
async def resolve_image_node(
root: Any,
info: graphene.ResolveInfo,
id: ResolvedGlobalID,
scope_id: ScopeType,
permission: ImagePermission = ImagePermission.READ_ATTRIBUTE,
) -> Optional[ImageNode]:
return await ImageNode.get_node(info, id, scope_id, permission)

@staticmethod
async def resolve_image_nodes(
root: Any,
info: graphene.ResolveInfo,
*,
scope_id: ScopeType,
permission: ImagePermission = ImagePermission.READ_ATTRIBUTE,
filter: Optional[str] = None,
order: Optional[str] = None,
offset: Optional[int] = None,
after: Optional[str] = None,
first: Optional[int] = None,
before: Optional[str] = None,
last: Optional[int] = None,
) -> ConnectionResolverResult[ImageNode]:
return await ImageNode.get_connection(
info,
scope_id,
permission,
filter_expr=filter,
order_expr=order,
offset=offset,
after=after,
first=first,
before=before,
last=last,
)

@staticmethod
@scoped_query(autofill_user=True, user_key="access_key")
async def resolve_keypair(
Expand Down
Loading

0 comments on commit b14387f

Please sign in to comment.