feat: Implement CRUD API for managing Harbor per-project Quota

changes/3090.feature.md

@@ -0,0 +1 @@
+Implement CRUD API for managing Harbor per-project Quota.

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

@@ -3,6 +3,7 @@
 
 from ai.backend.client.output.fields import group_fields
 from ai.backend.client.output.types import FieldSpec
+from ai.backend.common.utils import b64encode
 
 from ...cli.types import Undefined, undefined
 from ..session import api_session
@@ -311,3 +312,101 @@ async def remove_users(
         }
         data = await api_session.get().Admin._query(query, variables)
         return data["modify_group"]
+
+    @api_function
+    @classmethod
+    async def get_container_registry_quota(cls, group_id: str) -> int:
+        """
+        Get Quota Limit for the group's container registry.
+        Currently only HarborV2 registry is supported.
+
+        You need an admin privilege for this operation.
+        """
+        query = textwrap.dedent(
+            """\
+            query($id: String!) {
+                group_node(id: $id) {
+                    registry_quota
+                }
+            }
+        """
+        )
+
+        variables = {"id": b64encode(f"group_node:{group_id}")}
+        data = await api_session.get().Admin._query(query, variables)
+        return data["group_node"]["registry_quota"]
+
+    @api_function
+    @classmethod
+    async def create_container_registry_quota(cls, group_id: str, quota: int) -> dict:
+        """
+        Create Quota Limit for the group's container registry.
+        Currently only HarborV2 registry is supported.
+
+        You need an admin privilege for this operation.
+        """
+        query = textwrap.dedent(
+            """\
+            mutation($scope_id: ScopeField!, $quota: Int!) {
+                create_container_registry_quota(
+                        scope_id: $scope_id, quota: $quota) {
+                    ok msg
+                }
+            }
+        """
+        )
+
+        scope_id = f"project:{group_id}"
+        variables = {"scope_id": scope_id, "quota": quota}
+        data = await api_session.get().Admin._query(query, variables)
+        return data["create_container_registry_quota"]
+
+    @api_function
+    @classmethod
+    async def update_container_registry_quota(cls, group_id: str, quota: int) -> dict:
+        """
+        Update Quota Limit for the group's container registry.
+        Currently only HarborV2 registry is supported.
+
+        You need an admin privilege for this operation.
+        """
+        query = textwrap.dedent(
+            """\
+            mutation($scope_id: ScopeField!, $quota: Int!) {
+                update_container_registry_quota(
+                        scope_id: $scope_id, quota: $quota) {
+                    ok msg
+                }
+            }
+        """
+        )
+
+        scope_id = f"project:{group_id}"
+        variables = {"scope_id": scope_id, "quota": quota}
+        data = await api_session.get().Admin._query(query, variables)
+        return data["update_container_registry_quota"]
+
+    @api_function
+    @classmethod
+    async def delete_container_registry_quota(cls, group_id: str) -> dict:
+        """
+        Delete Quota Limit for the group's container registry.
+        Currently only HarborV2 registry is supported.
+
+        You need an admin privilege for this operation.
+        """
+        query = textwrap.dedent(
+            """\
+            mutation($scope_id: ScopeField!) {
+                delete_container_registry_quota(
+                        scope_id: $scope_id) {
+                    ok msg
+                }
+            }
+        """
+        )
+
+        scope_id = f"project:{group_id}"
+        variables = {"scope_id": scope_id}
+        data = await api_session.get().Admin._query(query, variables)
+        return data["delete_container_registry_quota"]

src/ai/backend/common/utils.py

@@ -425,3 +425,12 @@ def join_non_empty(*args, sep):
     """
     filtered_args = [arg for arg in args if arg]
     return sep.join(filtered_args)
+
+
+def b64encode(s: str) -> str:
+    """
+    base64 encoding method of graphql_relay.
+    Use it in components where the graphql_relay package is unavailable.
+    """
+    b: bytes = s.encode("utf-8") if isinstance(s, str) else s
+    return base64.b64encode(b).decode("ascii")

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/group.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Iterable, Tuple
+
+import aiohttp_cors
+import trafaret as t
+from aiohttp import web
+
+from ai.backend.common import validators as tx
+from ai.backend.logging import BraceStyleAdapter
+from ai.backend.manager.models.gql_models.container_registry_utils import (
+    HarborQuotaManager,
+)
+from ai.backend.manager.models.rbac import ProjectScope
+
+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 check_api_params
+
+log = BraceStyleAdapter(logging.getLogger(__spec__.name))
+
+
+@server_status_required(READ_ALLOWED)
+@superadmin_required
+@check_api_params(
+    t.Dict({
+        tx.AliasedKey(["group_id", "group"]): t.String,
+        tx.AliasedKey(["quota"]): t.Int,
+    })
+)
+async def update_registry_quota(request: web.Request, params: Any) -> web.Response:
+    log.info("UPDATE_REGISTRY_QUOTA (gr:{})", params["group_id"])
+    root_ctx: RootContext = request.app["_root.context"]
+    group_id = params["group_id"]
+    scope_id = ProjectScope(project_id=group_id, domain_name=None)
+    quota = int(params["quota"])
+
+    async with root_ctx.db.begin_session() as db_sess:
+        manager = await HarborQuotaManager.new(db_sess, scope_id)
+        await manager.update(quota)
+
+    return web.json_response({})
+
+
+@server_status_required(READ_ALLOWED)
+@superadmin_required
+@check_api_params(
+    t.Dict({
+        tx.AliasedKey(["group_id", "group"]): t.String,
+    })
+)
+async def delete_registry_quota(request: web.Request, params: Any) -> web.Response:
+    log.info("DELETE_REGISTRY_QUOTA (gr:{})", params["group_id"])
+    root_ctx: RootContext = request.app["_root.context"]
+    group_id = params["group_id"]
+    scope_id = ProjectScope(project_id=group_id, domain_name=None)
+
+    async with root_ctx.db.begin_session() as db_sess:
+        manager = await HarborQuotaManager.new(db_sess, scope_id)
+        await manager.delete()
+
+    return web.json_response({})
+
+
+@server_status_required(READ_ALLOWED)
+@superadmin_required
+@check_api_params(
+    t.Dict({
+        tx.AliasedKey(["group_id", "group"]): t.String,
+        tx.AliasedKey(["quota"]): t.Int,
+    })
+)
+async def create_registry_quota(request: web.Request, params: Any) -> web.Response:
+    log.info("CREATE_REGISTRY_QUOTA (gr:{})", params["group_id"])
+    root_ctx: RootContext = request.app["_root.context"]
+    group_id = params["group_id"]
+    scope_id = ProjectScope(project_id=group_id, domain_name=None)
+    quota = int(params["quota"])
+
+    async with root_ctx.db.begin_session() as db_sess:
+        manager = await HarborQuotaManager.new(db_sess, scope_id)
+        await manager.create(quota)
+
+    return web.json_response({})
+
+
+@server_status_required(READ_ALLOWED)
+@superadmin_required
+@check_api_params(
+    t.Dict({
+        tx.AliasedKey(["group_id", "group"]): t.String,
+    })
+)
+async def read_registry_quota(request: web.Request, params: Any) -> web.Response:
+    log.info("READ_REGISTRY_QUOTA (gr:{})", params["group_id"])
+    root_ctx: RootContext = request.app["_root.context"]
+    group_id = params["group_id"]
+    scope_id = ProjectScope(project_id=group_id, domain_name=None)
+
+    async with root_ctx.db.begin_session() as db_sess:
+        manager = await HarborQuotaManager.new(db_sess, scope_id)
+        quota = await manager.read()
+
+    return web.json_response({"result": quota})
+
+
+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"] = "group"
+    cors = aiohttp_cors.setup(app, defaults=default_cors_options)
+    cors.add(app.router.add_route("POST", "/registry-quota", create_registry_quota))
+    cors.add(app.router.add_route("GET", "/registry-quota", read_registry_quota))
+    cors.add(app.router.add_route("PATCH", "/registry-quota", update_registry_quota))
+    cors.add(app.router.add_route("DELETE", "/registry-quota", delete_registry_quota))
+    return app, []

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

@@ -698,6 +698,9 @@
   """Added in 24.03.7."""
   container_registry: JSONString
   scaling_groups: [String]
+
+  """Added in 24.12.0."""
+  registry_quota: BigInt
   user_nodes(filter: String, order: String, offset: Int, before: String, after: String, first: Int, last: Int): UserConnection
 }
 
@@ -1844,6 +1847,15 @@
 
   """Added in 24.12.0"""
   disassociate_container_registry_with_group(group_id: String!, registry_id: String!): DisassociateContainerRegistryWithGroup
+
+  """Added in 24.12.0"""
+  create_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): CreateContainerRegistryQuota
+
+  """Added in 24.12.0"""
+  update_container_registry_quota(quota: BigInt!, scope_id: ScopeField!): UpdateContainerRegistryQuota
+
+  """Added in 24.12.0"""
+  delete_container_registry_quota(scope_id: ScopeField!): DeleteContainerRegistryQuota
   create_container_registry(hostname: String!, props: CreateContainerRegistryInput!): CreateContainerRegistry
   modify_container_registry(hostname: String!, props: ModifyContainerRegistryInput!): ModifyContainerRegistry
   delete_container_registry(hostname: String!): DeleteContainerRegistry
@@ -2570,6 +2582,24 @@
   msg: String
 }
 
+"""Added in 24.12.0."""
+type CreateContainerRegistryQuota {
+  ok: Boolean
+  msg: String
+}
+
+"""Added in 24.12.0."""
+type UpdateContainerRegistryQuota {
+  ok: Boolean
+  msg: String
+}
+
+"""Added in 24.12.0."""
+type DeleteContainerRegistryQuota {
+  ok: Boolean
+  msg: String
+}
+
 type CreateContainerRegistry {
   container_registry: ContainerRegistry
 }

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

@@ -74,7 +74,10 @@
 )
 from .gql_models.container_registry import (
     AssociateContainerRegistryWithGroup,
+    CreateContainerRegistryQuota,
+    DeleteContainerRegistryQuota,
     DisassociateContainerRegistryWithGroup,
+    UpdateContainerRegistryQuota,
 )
 from .gql_models.domain import (
     CreateDomainNode,
@@ -349,6 +352,16 @@ class Mutations(graphene.ObjectType):
         description="Added in 24.12.0"
     )
 
+    create_container_registry_quota = CreateContainerRegistryQuota.Field(
+        description="Added in 24.12.0"
+    )
+    update_container_registry_quota = UpdateContainerRegistryQuota.Field(
+        description="Added in 24.12.0"
+    )
+    delete_container_registry_quota = DeleteContainerRegistryQuota.Field(
+        description="Added in 24.12.0"
+    )
+
     # Legacy mutations
     create_container_registry = CreateContainerRegistry.Field()
     modify_container_registry = ModifyContainerRegistry.Field()