Skip to content

Commit

Permalink
feat: Add registry_quota to GroupNode
Browse files Browse the repository at this point in the history
  • Loading branch information
jopemachine committed Nov 13, 2024
1 parent 179bf16 commit 08ff80e
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 0 deletions.
3 changes: 3 additions & 0 deletions src/ai/backend/manager/api/schema.graphql
Original file line number Diff line number Diff line change
Expand Up @@ -689,6 +689,9 @@ type GroupNode implements Node {
"""Added in 24.03.7."""
container_registry: JSONString
scaling_groups: [String]

"""Added in 24.12.0."""
registry_quota: Int
user_nodes(filter: String, order: String, offset: Int, before: String, after: String, first: Int, last: Int): UserConnection
}

Expand Down
68 changes: 68 additions & 0 deletions src/ai/backend/manager/models/gql_models/group.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,20 @@
from collections.abc import Mapping
from typing import (
TYPE_CHECKING,
Any,
Self,
Sequence,
)

import aiohttp
import graphene
import sqlalchemy as sa
import yarl
from dateutil.parser import parse as dtparse
from graphene.types.datetime import DateTime as GQLDateTime

from ai.backend.manager.models.container_registry import ContainerRegistryRow, ContainerRegistryType

from ..base import (
FilterExprArg,
OrderExprArg,
Expand Down Expand Up @@ -112,6 +117,8 @@ class Meta:
lambda: graphene.String,
)

registry_quota = graphene.Int(description="Added in 24.12.0.")

user_nodes = PaginatedConnectionField(
UserConnection,
)
Expand Down Expand Up @@ -204,6 +211,67 @@ async def resolve_user_nodes(
total_cnt = await db_session.scalar(cnt_query)
return ConnectionResolverResult(result, cursor, pagination_order, page_size, total_cnt)

async def resolve_registry_quota(self, info: graphene.ResolveInfo) -> int:
graph_ctx = info.context

# user = graph_ctx.user
# client_ctx = ClientContext(
# graph_ctx.db, user["domain_name"], user["uuid"], user["role"]
# )

async with graph_ctx.db.begin_session() as db_sess:
if (
self.container_registry is None
or "registry" not in self.container_registry
or "project" not in self.container_registry
):
raise ValueError("Container registry info does not exist in the group.")

registry_name, project = (
self.container_registry["registry"],
self.container_registry["project"],
)

cr_query = sa.select(ContainerRegistryRow).where(
(ContainerRegistryRow.registry_name == registry_name)
& (ContainerRegistryRow.project == project)
)

result = await db_sess.execute(cr_query)
registry = result.fetchone()[0]

if registry.type != ContainerRegistryType.HARBOR2:
raise ValueError("Only HarborV2 registry is supported for now.")

ssl_verify = registry.ssl_verify
connector = aiohttp.TCPConnector(ssl=ssl_verify)

api_url = yarl.URL(registry.url) / "api" / "v2.0"
async with aiohttp.ClientSession(connector=connector) as sess:
rqst_args: dict[str, Any] = {}
rqst_args["auth"] = aiohttp.BasicAuth(
registry.username,
registry.password,
)

get_project_id_api = api_url / "projects" / project

async with sess.get(get_project_id_api, allow_redirects=False, **rqst_args) as resp:
res = await resp.json()
harbor_project_id = res["project_id"]

get_quota_id_api = (api_url / "quotas").with_query({
"reference": "project",
"reference_id": harbor_project_id,
})

async with sess.get(get_quota_id_api, allow_redirects=False, **rqst_args) as resp:
res = await resp.json()
# TODO: Raise error when quota is not found or multiple quotas are found.
quota = res[0]["hard"]["storage"]

return quota

@classmethod
async def get_node(cls, info: graphene.ResolveInfo, id) -> Self:
graph_ctx: GraphQueryContext = info.context
Expand Down

0 comments on commit 08ff80e

Please sign in to comment.