fix: Update GQLPrivilegeCheckMiddleware to align with upstream chan…

…ges on `graphql-core` package (#2598) Co-authored-by: Joongi Kim <[email protected]> Co-authored-by: Jihyun Kang <[email protected]>
lablup · Jul 31, 2024 · 187cf9e · 187cf9e
1 parent fb196c8
commit 187cf9e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/changes/2598.fix.md b/changes/2598.fix.md
@@ -0,0 +1 @@
+Update `GQLPrivilegeCheckMiddleware` to align with upstream changes on `graphql-core` package
diff --git a/src/ai/backend/manager/models/gql.py b/src/ai/backend/manager/models/gql.py
@@ -7,7 +7,8 @@
 import attrs
 import graphene
 from graphene.types.inputobjecttype import set_input_object_type_default_value
-from graphql import Undefined
+from graphql import OperationType, Undefined
+from graphql.type import GraphQLField
 
 set_input_object_type_default_value(Undefined)
 
@@ -2240,10 +2241,13 @@ async def resolve_model_cards(
 class GQLMutationPrivilegeCheckMiddleware:
     def resolve(self, next, root, info: graphene.ResolveInfo, **args) -> Any:
         graph_ctx: GraphQueryContext = info.context
-        if info.operation.operation == "mutation" and len(info.path) == 1:
-            mutation_cls = getattr(Mutations, info.field_name).type
+        if info.operation.operation == OperationType.MUTATION:
+            mutation_field: GraphQLField | None = getattr(Mutations, info.field_name, None)  # noqa
+            if mutation_field is None:
+                return next(root, info, **args)
+            mutation_cls = mutation_field.type
             # default is allow nobody.
             allowed_roles = getattr(mutation_cls, "allowed_roles", [])
             if graph_ctx.user["role"] not in allowed_roles:
-                return mutation_cls(False, f"no permission to execute {info.path[0]}")
+                return mutation_cls(False, f"no permission to execute {info.path.key}")  # type: ignore
         return next(root, info, **args)