From 739482141c6fa1ba68a3e23af6ed5cd4baf15b1e Mon Sep 17 00:00:00 2001
From: jopemachine <jopemachine@naver.com>
Date: Tue, 11 Feb 2025 04:49:34 +0000
Subject: [PATCH] fix: Remove ok, msg from PurgeImage mutation

---
 docs/manager/graphql-reference/schema.graphql     |  5 +----
 src/ai/backend/manager/models/gql_models/image.py | 14 ++++++--------
 2 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/docs/manager/graphql-reference/schema.graphql b/docs/manager/graphql-reference/schema.graphql
index cc09f11b3f7..939cdc38deb 100644
--- a/docs/manager/graphql-reference/schema.graphql
+++ b/docs/manager/graphql-reference/schema.graphql
@@ -2362,14 +2362,11 @@ type ForgetImage {
 
 """Added in 25.3.0."""
 type PurgeImageById {
-  ok: Boolean
-  msg: String
   image: ImageNode
 }
 
+"""Added in 25.3.0."""
 type PurgeImage {
-  ok: Boolean
-  msg: String
   image: ImageNode
 }
 
diff --git a/src/ai/backend/manager/models/gql_models/image.py b/src/ai/backend/manager/models/gql_models/image.py
index caa7e94f105..9a29f072f93 100644
--- a/src/ai/backend/manager/models/gql_models/image.py
+++ b/src/ai/backend/manager/models/gql_models/image.py
@@ -29,7 +29,7 @@
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.manager.models.container_registry import ContainerRegistryRow, ContainerRegistryType
 
-from ...api.exceptions import ImageNotFound, ObjectNotFound
+from ...api.exceptions import GenericForbidden, ImageNotFound, ObjectNotFound
 from ...defs import DEFAULT_IMAGE_ARCH
 from ..base import batch_multiresult_in_scalar_stream, set_if_set
 from ..gql_relay import AsyncNode
@@ -568,6 +568,8 @@ async def mutate(
 
 
 class PurgeImage(graphene.Mutation):
+    """Added in 25.3.0."""
+
     allowed_roles = (
         UserRole.SUPERADMIN,
         UserRole.ADMIN,
@@ -578,8 +580,6 @@ class Arguments:
         reference = graphene.String(required=True)
         architecture = graphene.String(default_value=DEFAULT_IMAGE_ARCH)
 
-    ok = graphene.Boolean()
-    msg = graphene.String()
     image = graphene.Field(ImageNode)
 
     @staticmethod
@@ -611,7 +611,7 @@ async def mutate(
                 ):
                     return PurgeImage(ok=False, msg="Forbidden")
             await session.delete(image_row)
-            return PurgeImage(ok=True, msg="", image=ImageNode.from_row(image_row))
+            return PurgeImage(image=ImageNode.from_row(image_row))
 
 
 class PurgeImageById(graphene.Mutation):
@@ -626,8 +626,6 @@ class PurgeImageById(graphene.Mutation):
     class Arguments:
         image_id = graphene.String(required=True)
 
-    ok = graphene.Boolean()
-    msg = graphene.String()
     image = graphene.Field(ImageNode)
 
     @staticmethod
@@ -661,9 +659,9 @@ async def mutate(
                     not customized_image_owner
                     or customized_image_owner != f"user:{ctx.user['uuid']}"
                 ):
-                    return PurgeImageById(ok=False, msg="Forbidden")
+                    raise GenericForbidden("Image is not owned by your account.")
             await session.delete(image_row)
-            return PurgeImageById(ok=True, msg="", image=ImageNode.from_row(image_row))
+            return PurgeImageById(image=ImageNode.from_row(image_row))
 
 
 class UntagImageFromRegistry(graphene.Mutation):