Skip to content

Commit

Permalink
fix: Remove useless strenum
Browse files Browse the repository at this point in the history
  • Loading branch information
jopemachine committed Nov 14, 2024
1 parent f8a599d commit 33c1b79
Showing 1 changed file with 10 additions and 20 deletions.
30 changes: 10 additions & 20 deletions src/ai/backend/manager/models/gql_models/container_registry.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
from __future__ import annotations

import enum
import logging
from typing import Any, Self
from typing import Any, Literal, Self

import aiohttp
import aiohttp.client_exceptions
Expand Down Expand Up @@ -81,19 +80,16 @@ async def mutate(
return await simple_db_mutate(cls, info.context, delete_query)


class UpdateQuotaOperationType(enum.StrEnum):
CREATE = "create"
DELETE = "delete"
UPDATE = "update"


async def update_harbor_project_quota(
operation_type: UpdateQuotaOperationType,
operation_type: Literal["create", "delete", "update"],
mutation_cls: Any,
info: graphene.ResolveInfo,
scope_id: ScopeType,
quota: int,
) -> Any:
"""
Utility function for code reuse of the HarborV2 per-project Quota CRUD API
"""
if not isinstance(scope_id, ProjectScope):
return mutation_cls(
ok=False, msg="Quota mutation currently supports only the project scope."
Expand Down Expand Up @@ -178,10 +174,10 @@ async def update_harbor_project_quota(
)

previous_quota = res[0]["hard"]["storage"]
if operation_type == UpdateQuotaOperationType.DELETE:
if operation_type == "delete":
if previous_quota == -1:
return mutation_cls(ok=False, msg=f"Quota entity not found. (gr: {project_id})")
elif operation_type == UpdateQuotaOperationType.CREATE:
elif operation_type == "create":
if previous_quota > 0:
return mutation_cls(
ok=False, msg=f"Quota limit already exists. (gr: {project_id})"
Expand Down Expand Up @@ -229,9 +225,7 @@ async def mutate(
scope_id: ScopeType,
quota: int,
) -> Self:
return await update_harbor_project_quota(
UpdateQuotaOperationType.CREATE, cls, info, scope_id, quota
)
return await update_harbor_project_quota("create", cls, info, scope_id, quota)


class UpdateQuota(graphene.Mutation):
Expand All @@ -257,9 +251,7 @@ async def mutate(
scope_id: ScopeType,
quota: int,
) -> Self:
return await update_harbor_project_quota(
UpdateQuotaOperationType.UPDATE, cls, info, scope_id, quota
)
return await update_harbor_project_quota("update", cls, info, scope_id, quota)


class DeleteQuota(graphene.Mutation):
Expand All @@ -283,6 +275,4 @@ async def mutate(
info: graphene.ResolveInfo,
scope_id: ScopeType,
) -> Self:
return await update_harbor_project_quota(
UpdateQuotaOperationType.DELETE, cls, info, scope_id, -1
)
return await update_harbor_project_quota("delete", cls, info, scope_id, -1)

0 comments on commit 33c1b79

Please sign in to comment.