From 0d5ccf272cc95cef43d1928ccdad87183fb8a5f4 Mon Sep 17 00:00:00 2001
From: Gyubong Lee <jopemachine@naver.com>
Date: Thu, 14 Nov 2024 04:08:21 +0000
Subject: [PATCH] fix: Remove useless strenum

---
 .../models/gql_models/container_registry.py   | 30 +++++++------------
 1 file changed, 10 insertions(+), 20 deletions(-)

diff --git a/src/ai/backend/manager/models/gql_models/container_registry.py b/src/ai/backend/manager/models/gql_models/container_registry.py
index b27dc7c5f2..bf75cef10b 100644
--- a/src/ai/backend/manager/models/gql_models/container_registry.py
+++ b/src/ai/backend/manager/models/gql_models/container_registry.py
@@ -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
@@ -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."
@@ -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})"
@@ -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):
@@ -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):
@@ -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)