fix: DeferStreamDirectiveOnRootField forbids fragments on abstract types

Replicates graphql/graphql-js@3aad6a3

Author: Cito

src/graphql/validation/rules/defer_stream_directive_on_root_field.py

@@ -2,72 +2,121 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any
 
 from ...error import GraphQLError
+from ...language import (
+    FieldNode,
+    FragmentDefinitionNode,
+    FragmentSpreadNode,
+    InlineFragmentNode,
+    OperationType,
+)
 from ...type import GraphQLDeferDirective, GraphQLStreamDirective
-from . import ASTValidationRule, ValidationContext
+from . import ValidationRule
 
 if TYPE_CHECKING:
-    from ...language import DirectiveNode, Node
+    from ...language import (
+        DirectiveNode,
+        OperationDefinitionNode,
+        SelectionSetNode,
+    )
+    from ...type import GraphQLObjectType
 
 __all__ = ["DeferStreamDirectiveOnRootField"]
 
 
-class DeferStreamDirectiveOnRootField(ASTValidationRule):
+def get_directive(
+    node: FieldNode | FragmentSpreadNode | InlineFragmentNode, name: str
+) -> DirectiveNode | None:
+    for directive in node.directives or ():
+        if directive.name.value == name:
+            return directive
+    return None
+
+
+class DeferStreamDirectiveOnRootField(ValidationRule):
     """Defer and stream directives are used on valid root field
 
     A GraphQL document is only valid if defer directives are not used on root
     mutation or subscription types.
     """
 
-    def enter_directive(
-        self,
-        node: DirectiveNode,
-        _key: Any,
-        _parent: Any,
-        _path: Any,
-        _ancestors: list[Node],
+    def enter_operation_definition(
+        self, node: OperationDefinitionNode, *_args: Any
     ) -> None:
-        context = cast("ValidationContext", self.context)
-        parent_type = context.get_parent_type()
-        if not parent_type:
+        operation = node.operation
+        if operation not in (OperationType.SUBSCRIPTION, OperationType.MUTATION):
             return
-        schema = context.schema
-        mutation_type = schema.mutation_type
-        subscription_type = schema.subscription_type
+        schema = self.context.schema
+        root_type = schema.get_root_type(operation)
+        if root_type:
+            document = self.context.document
+            fragments = {
+                definition.name.value: definition
+                for definition in document.definitions
+                if isinstance(definition, FragmentDefinitionNode)
+            }
+            self.forbid_defer_stream(
+                operation, root_type, fragments, node.selection_set, set()
+            )
 
-        if node.name.value == GraphQLDeferDirective.name:
-            if mutation_type and parent_type is mutation_type:
-                self.report_error(
-                    GraphQLError(
-                        "Defer directive cannot be used on root"
-                        f" mutation type '{parent_type.name}'.",
-                        node,
-                    )
-                )
-            if subscription_type and parent_type is subscription_type:
-                self.report_error(
-                    GraphQLError(
-                        "Defer directive cannot be used on root"
-                        f" subscription type '{parent_type.name}'.",
-                        node,
+    def forbid_defer_stream(
+        self,
+        operation_type: OperationType,
+        root_type: GraphQLObjectType,
+        fragments: dict[str, FragmentDefinitionNode],
+        selection_set: SelectionSetNode,
+        visited_fragments: set[str],
+    ) -> None:
+        for selection in selection_set.selections:
+            if isinstance(selection, FieldNode):
+                stream = get_directive(selection, GraphQLStreamDirective.name)
+                if stream:
+                    self.report_error(
+                        GraphQLError(
+                            "Stream directive cannot be used on root"
+                            f" {operation_type.value} type '{root_type.name}'.",
+                            stream,
+                        )
                     )
-                )
-        if node.name.value == GraphQLStreamDirective.name:
-            if mutation_type and parent_type is mutation_type:
-                self.report_error(
-                    GraphQLError(
-                        "Stream directive cannot be used on root"
-                        f" mutation type '{parent_type.name}'.",
-                        node,
+            elif isinstance(selection, FragmentSpreadNode):
+                fragment_name = selection.name.value
+                if fragment_name in visited_fragments:
+                    continue
+                fragment = fragments.get(fragment_name)
+                if fragment:
+                    defer = get_directive(selection, GraphQLDeferDirective.name)
+                    if defer:
+                        self.report_error(
+                            GraphQLError(
+                                "Defer directive cannot be used on root"
+                                f" {operation_type.value} type '{root_type.name}'.",
+                                defer,
+                            )
+                        )
+                    self.forbid_defer_stream(
+                        operation_type,
+                        root_type,
+                        fragments,
+                        fragment.selection_set,
+                        visited_fragments,
                     )
-                )
-            if subscription_type and parent_type is subscription_type:
-                self.report_error(
-                    GraphQLError(
-                        "Stream directive cannot be used on root"
-                        f" subscription type '{parent_type.name}'.",
-                        node,
+                visited_fragments.add(fragment_name)
+            elif isinstance(selection, InlineFragmentNode):
+                defer = get_directive(selection, GraphQLDeferDirective.name)
+                if defer:
+                    self.report_error(
+                        GraphQLError(
+                            "Defer directive cannot be used on root"
+                            f" {operation_type.value} type '{root_type.name}'.",
+                            defer,
+                        )
                     )
+                self.forbid_defer_stream(
+                    operation_type,
+                    root_type,
+                    fragments,
+                    selection.selection_set,
+                    visited_fragments,
                 )

tests/validation/test_defer_stream_directive_on_root_field.py

@@ -12,19 +12,26 @@
       sender: String
     }
 
-    type SubscriptionRoot {
+    interface Root {
+      rootField: Message
+    }
+
+    type SubscriptionRoot implements Root {
       subscriptionField: Message
       subscriptionListField: [Message]
+      rootField: Message
     }
 
-    type MutationRoot {
+    type MutationRoot implements Root {
       mutationField: Message
       mutationListField: [Message]
+      rootField: Message
     }
 
-    type QueryRoot {
+    type QueryRoot implements Root {
       message: Message
       messages: [Message]
+      rootField: Message
     }
 
     schema {
@@ -75,6 +82,13 @@ def defer_fragment_spread_on_root_mutation_field():
             """
             mutation {
               ...rootFragment @defer
+              ...otherFragment
+            }
+            fragment otherFragment on MutationRoot {
+              ...rootFragment
+              mutationListField {
+                body
+              }
             }
             fragment rootFragment on MutationRoot {
               mutationField {
@@ -91,6 +105,29 @@ def defer_fragment_spread_on_root_mutation_field():
             ],
         )
 
+    def defer_fragment_spread_on_root_mutation_field_interface():
+        assert_errors(
+            """
+            mutation {
+              ...rootFragment
+            }
+            fragment rootFragment on Root {
+              ... @defer {
+                rootField {
+                  body
+                }
+              }
+            }
+            """,
+            [
+                {
+                    "message": "Defer directive cannot be used on root"
+                    " mutation type 'MutationRoot'.",
+                    "locations": [(6, 19)],
+                },
+            ],
+        )
+
     def defer_inline_fragment_spread_on_root_mutation_field():
         assert_errors(
             """
@@ -124,6 +161,29 @@ def defer_fragment_spread_on_nested_mutation_field():
             """
         )
 
+    def defer_fragment_spread_on_root_subscription_field_interface():
+        assert_errors(
+            """
+            subscription {
+              ...rootFragment
+            }
+            fragment rootFragment on Root {
+              ... @defer {
+                rootField {
+                  body
+                }
+              }
+            }
+            """,
+            [
+                {
+                    "message": "Defer directive cannot be used on root"
+                    " subscription type 'SubscriptionRoot'.",
+                    "locations": [(6, 19)],
+                },
+            ],
+        )
+
     def defer_fragment_spread_on_root_subscription_field():
         assert_errors(
             """