feat(BA-463): Implement APIs for associating, disassociating container_registries with groups

changes/3067.feature.md

@@ -0,0 +1 @@
+Implement APIs for associating, disassociating `container_registries` with `groups`.

src/ai/backend/client/func/container_registry.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import textwrap
+
+from ..session import api_session
+from .base import BaseFunction, api_function
+
+__all__ = ("ContainerRegistry",)
+
+
+class ContainerRegistry(BaseFunction):
+    """
+    Provides a shortcut of :func:`Admin.query()
+    <ai.backend.client.admin.Admin.query>` that fetches, modifies various container registry
+    information.
+
+    .. note::
+
+      All methods in this function class require your API access key to
+      have the *admin* privilege.
+    """
+
+    @api_function
+    @classmethod
+    async def associate_group(cls, registry_id: str, group_id: str) -> dict:
+        """
+        Associate container_registry with group.
+
+        :param registry_id: ID of the container registry.
+        :param group_id: ID of the group.
+        """
+        query = textwrap.dedent(
+            """\
+            mutation($registry_id: String!, $group_id: String!) {
+                associate_container_registry_with_group(
+                        registry_id: $registry_id, group_id: $group_id) {
+                    ok msg
+                }
+            }
+        """
+        )
+        variables = {"registry_id": registry_id, "group_id": group_id}
+        data = await api_session.get().Admin._query(query, variables)
+        return data["associate_container_registry_with_group"]
+
+    @api_function
+    @classmethod
+    async def disassociate_group(cls, registry_id: str, group_id: str) -> dict:
+        """
+        Disassociate container_registry with group.
+
+        :param registry_id: ID of the container registry.
+        :param group_id: ID of the group.
+        """
+        query = textwrap.dedent(
+            """\
+            mutation($registry_id: String!, $group_id: String!) {
+                disassociate_container_registry_with_group(
+                        registry_id: $registry_id, group_id: $group_id) {
+                    ok msg
+                }
+            }
+        """
+        )
+        variables = {"registry_id": registry_id, "group_id": group_id}
+        data = await api_session.get().Admin._query(query, variables)
+        return data["disassociate_container_registry_with_group"]

src/ai/backend/client/session.py

@@ -254,6 +254,7 @@ class BaseSession(metaclass=abc.ABCMeta):
         "ScalingGroup",
         "Storage",
         "Image",
+        "ContainerRegistry",
         "ComputeSession",
         "SessionTemplate",
         "Domain",
@@ -298,6 +299,7 @@ def __init__(
         from .func.agent import Agent, AgentWatcher
         from .func.auth import Auth
         from .func.bgtask import BackgroundTask
+        from .func.container_registry import ContainerRegistry
         from .func.domain import Domain
         from .func.dotfile import Dotfile
         from .func.etcd import EtcdConfig
@@ -327,6 +329,7 @@ def __init__(
         self.Storage = Storage
         self.Auth = Auth
         self.BackgroundTask = BackgroundTask
+        self.ContainerRegistry = ContainerRegistry
         self.EtcdConfig = EtcdConfig
         self.Domain = Domain
         self.Group = Group

src/ai/backend/manager/api/container_registry.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Iterable, Tuple
+
+import aiohttp_cors
+import sqlalchemy as sa
+from aiohttp import web
+from pydantic import AliasChoices, BaseModel, Field
+from sqlalchemy.exc import IntegrityError
+
+from ai.backend.logging import BraceStyleAdapter
+from ai.backend.manager.models.association_container_registries_groups import (
+    AssociationContainerRegistriesGroupsRow,
+)
+
+from .exceptions import ContainerRegistryNotFound, GenericBadRequest
+
+if TYPE_CHECKING:
+    from .context import RootContext
+
+from .auth import superadmin_required
+from .manager import READ_ALLOWED, server_status_required
+from .types import CORSOptions, WebMiddleware
+from .utils import pydantic_params_api_handler
+
+log = BraceStyleAdapter(logging.getLogger(__spec__.name))
+
+
+class AssociationRequestModel(BaseModel):
+    registry_id: str = Field(
+        validation_alias=AliasChoices("registry_id", "registry"),
+        description="Container registry row's ID",
+    )
+    group_id: str = Field(
+        validation_alias=AliasChoices("group_id", "group"),
+        description="Group row's ID",
+    )
+
+
+@server_status_required(READ_ALLOWED)
+@superadmin_required
+@pydantic_params_api_handler(AssociationRequestModel)
+async def associate_with_group(
+    request: web.Request, params: AssociationRequestModel
+) -> web.Response:
+    log.info("ASSOCIATE_WITH_GROUP (cr:{}, gr:{})", params.registry_id, params.group_id)
+    root_ctx: RootContext = request.app["_root.context"]
+    registry_id = params.registry_id
+    group_id = params.group_id
+
+    async with root_ctx.db.begin_session() as db_sess:
+        insert_query = sa.insert(AssociationContainerRegistriesGroupsRow).values({
+            "registry_id": registry_id,
+            "group_id": group_id,
+        })
+
+        try:
+            await db_sess.execute(insert_query)
+        except IntegrityError:
+            raise GenericBadRequest("Association already exists.")
+
+    return web.Response(status=204)
+
+
+class DisassociationRequestModel(BaseModel):
+    registry_id: str = Field(
+        validation_alias=AliasChoices("registry_id", "registry"),
+        description="Container registry row's ID",
+    )
+    group_id: str = Field(
+        validation_alias=AliasChoices("group_id", "group"),
+        description="Group row's ID",
+    )
+
+
+@server_status_required(READ_ALLOWED)
+@superadmin_required
+@pydantic_params_api_handler(DisassociationRequestModel)
+async def disassociate_with_group(
+    request: web.Request, params: DisassociationRequestModel
+) -> web.Response:
+    log.info("DISASSOCIATE_WITH_GROUP (cr:{}, gr:{})", params.registry_id, params.group_id)
+    root_ctx: RootContext = request.app["_root.context"]
+    registry_id = params.registry_id
+    group_id = params.group_id
+
+    async with root_ctx.db.begin_session() as db_sess:
+        delete_query = (
+            sa.delete(AssociationContainerRegistriesGroupsRow)
+            .where(AssociationContainerRegistriesGroupsRow.registry_id == registry_id)
+            .where(AssociationContainerRegistriesGroupsRow.group_id == group_id)
+        )
+
+        result = await db_sess.execute(delete_query)
+        if result.rowcount == 0:
+            raise ContainerRegistryNotFound()
+
+    return web.Response(status=204)
+
+
+def create_app(
+    default_cors_options: CORSOptions,
+) -> Tuple[web.Application, Iterable[WebMiddleware]]:
+    app = web.Application()
+    app["api_versions"] = (1, 2, 3, 4, 5)
+    app["prefix"] = "container-registries"
+    cors = aiohttp_cors.setup(app, defaults=default_cors_options)
+    cors.add(app.router.add_route("POST", "/associate-with-group", associate_with_group))
+    cors.add(app.router.add_route("POST", "/disassociate-with-group", disassociate_with_group))
+    return app, []

src/ai/backend/manager/api/exceptions.py

@@ -242,6 +242,10 @@ class EndpointTokenNotFound(ObjectNotFound):
     object_name = "endpoint_token"
 
 
+class ContainerRegistryNotFound(ObjectNotFound):
+    object_name = "container_registry"
+
+
 class TooManySessionsMatched(BackendError, web.HTTPNotFound):
     error_type = "https://api.backend.ai/probs/too-many-sessions-matched"
     error_title = "Too many sessions matched."

src/ai/backend/manager/api/schema.graphql

@@ -1,4 +1,4 @@
 schema {
   query: Queries
   mutation: Mutations
 }
@@ -1527,7 +1527,7 @@
   extra: JSONString
 }
 
 """Added in 24.09.0."""
 scalar ContainerRegistryTypeField
 
 """Added in 24.09.0."""
@@ -1655,7 +1655,7 @@
 type Mutations {
   modify_agent(id: String!, props: ModifyAgentInput!): ModifyAgent
   create_domain(name: String!, props: DomainInput!): CreateDomain
   modify_domain(name: String!, props: ModifyDomainInput!): ModifyDomain
 
   """Instead of deleting the domain, just mark it as inactive."""
   delete_domain(name: String!): DeleteDomain
@@ -1780,18 +1780,18 @@
   create_container_registry_node(
     """Added in 24.09.3."""
     extra: JSONString
 
     """Added in 24.09.0."""
     is_global: Boolean
 
     """Added in 24.09.0."""
     password: String
 
     """Added in 24.09.0."""
     project: String
 
     """Added in 24.09.0."""
     registry_name: String!
 
     """Added in 24.09.0."""
     ssl_verify: Boolean
@@ -1848,6 +1848,12 @@
     """Object id. Can be either global id or object id. Added in 24.09.0."""
     id: String!
   ): DeleteContainerRegistryNode
+
+  """Added in 25.2.0."""
+  associate_container_registry_with_group(group_id: String!, registry_id: String!): AssociateContainerRegistryWithGroup
+
+  """Added in 25.2.0."""
+  disassociate_container_registry_with_group(group_id: String!, registry_id: String!): DisassociateContainerRegistryWithGroup
   create_container_registry(hostname: String!, props: CreateContainerRegistryInput!): CreateContainerRegistry
   modify_container_registry(hostname: String!, props: ModifyContainerRegistryInput!): ModifyContainerRegistry
   delete_container_registry(hostname: String!): DeleteContainerRegistry
@@ -2516,7 +2522,7 @@
   msg: String
 }
 
 """Added in 24.03.9."""
 type DisassociateScalingGroupsWithDomain {
   ok: Boolean
   msg: String
@@ -2548,7 +2554,7 @@
   ok: Boolean
   msg: String
 }
 
 type DisassociateAllScalingGroupsWithGroup {
   ok: Boolean
   msg: String
@@ -2561,7 +2567,7 @@
 input QuotaScopeInput {
   hard_limit_bytes: BigInt
 }
 
 type UnsetQuotaScope {
   quota_scope: QuotaScope
 }
@@ -2581,6 +2587,18 @@
   container_registry: ContainerRegistryNode
 }
 
+"""Added in 25.2.0."""
+type AssociateContainerRegistryWithGroup {
+  ok: Boolean
+  msg: String
+}
+
+"""Added in 25.2.0."""
+type DisassociateContainerRegistryWithGroup {
+  ok: Boolean
+  msg: String
+}
+
 type CreateContainerRegistry {
   container_registry: ContainerRegistry
 }

src/ai/backend/manager/models/gql.py

@@ -74,6 +74,10 @@
     AgentSummaryList,
     ModifyAgent,
 )
+from .gql_models.container_registry import (
+    AssociateContainerRegistryWithGroup,
+    DisassociateContainerRegistryWithGroup,
+)
 from .gql_models.domain import (
     CreateDomainNode,
     DomainConnection,
@@ -335,6 +339,13 @@ class Mutations(graphene.ObjectType):
         description="Added in 24.09.0."
     )
 
+    associate_container_registry_with_group = AssociateContainerRegistryWithGroup.Field(
+        description="Added in 25.2.0."
+    )
+    disassociate_container_registry_with_group = DisassociateContainerRegistryWithGroup.Field(
+        description="Added in 25.2.0."
+    )
+
     # Legacy mutations
     create_container_registry = CreateContainerRegistry.Field()
     modify_container_registry = ModifyContainerRegistry.Field()

src/ai/backend/manager/models/gql_models/container_registry.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import logging
+from typing import Self
+
+import graphene
+import sqlalchemy as sa
+
+from ai.backend.logging import BraceStyleAdapter
+
+from ..association_container_registries_groups import (
+    AssociationContainerRegistriesGroupsRow,
+)
+from ..base import simple_db_mutate
+from ..user import UserRole
+
+log = BraceStyleAdapter(logging.getLogger(__spec__.name))  # type: ignore
+
+
+class AssociateContainerRegistryWithGroup(graphene.Mutation):
+    """Added in 25.2.0."""
+
+    allowed_roles = (UserRole.SUPERADMIN,)
+
+    class Arguments:
+        registry_id = graphene.String(required=True)
+        group_id = graphene.String(required=True)
+
+    ok = graphene.Boolean()
+    msg = graphene.String()
+
+    @classmethod
+    async def mutate(
+        cls,
+        root,
+        info: graphene.ResolveInfo,
+        registry_id: str,
+        group_id: str,
+    ) -> Self:
+        insert_query = sa.insert(AssociationContainerRegistriesGroupsRow).values({
+            "registry_id": registry_id,
+            "group_id": group_id,
+        })
+        return await simple_db_mutate(cls, info.context, insert_query)
+
+
+class DisassociateContainerRegistryWithGroup(graphene.Mutation):
+    """Added in 25.2.0."""
+
+    allowed_roles = (UserRole.SUPERADMIN,)
+
+    class Arguments:
+        registry_id = graphene.String(required=True)
+        group_id = graphene.String(required=True)
+
+    ok = graphene.Boolean()
+    msg = graphene.String()
+
+    @classmethod
+    async def mutate(
+        cls,
+        root,
+        info: graphene.ResolveInfo,
+        registry_id: str,
+        group_id: str,
+    ) -> Self:
+        delete_query = (
+            sa.delete(AssociationContainerRegistriesGroupsRow)
+            .where(AssociationContainerRegistriesGroupsRow.registry_id == registry_id)
+            .where(AssociationContainerRegistriesGroupsRow.group_id == group_id)
+        )
+        return await simple_db_mutate(cls, info.context, delete_query)

src/ai/backend/manager/server.py

@@ -170,6 +170,7 @@
 
 global_subapp_pkgs: Final[list[str]] = [
     ".acl",
+    ".container_registry",
     ".etcd",
     ".events",
     ".auth",