Rewrite DeferStreamDirectiveOnRootFieldRule to catch usage in fragments with abstract types

Author: robrichard

src/validation/__tests__/DeferStreamDirectiveOnRootFieldRule-test.ts

@@ -24,19 +24,26 @@ const schema = buildSchema(`
     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 {
@@ -76,6 +83,13 @@ describe('Validate: Defer/Stream directive on root field', () => {
     expectErrors(`
       mutation {
         ...rootFragment @defer
+        ...otherFragment
+      }
+      fragment otherFragment on MutationRoot {
+        ...rootFragment
+        mutationListField {
+          body
+        }
       }
       fragment rootFragment on MutationRoot {
         mutationField {
@@ -107,7 +121,26 @@ describe('Validate: Defer/Stream directive on root field', () => {
       },
     ]);
   });
-
+  it('Defer fragment spread on root mutation field interface', () => {
+    expectErrors(`
+      mutation {
+        ...rootFragment
+      }
+      fragment rootFragment on Root {
+        ... @defer {
+          rootField {
+            body
+          }
+        }
+      }
+    `).toDeepEqual([
+      {
+        message:
+          'Defer directive cannot be used on root mutation type "MutationRoot".',
+        locations: [{ line: 6, column: 13 }],
+      },
+    ]);
+  });
   it('Defer fragment spread on nested mutation field', () => {
     expectValid(`
       mutation {
@@ -120,6 +153,26 @@ describe('Validate: Defer/Stream directive on root field', () => {
     `);
   });
 
+  it('Defer fragment spread on root subscription field interface', () => {
+    expectErrors(`
+      subscription {
+        ...rootFragment
+      }
+      fragment rootFragment on Root {
+        ... @defer {
+            rootField {
+              body
+            }
+        }
+      }
+    `).toDeepEqual([
+      {
+        message:
+          'Defer directive cannot be used on root subscription type "SubscriptionRoot".',
+        locations: [{ line: 6, column: 13 }],
+      },
+    ]);
+  });
   it('Defer fragment spread on root subscription field', () => {
     expectErrors(`
       subscription {

src/validation/rules/DeferStreamDirectiveOnRootFieldRule.ts

@@ -1,11 +1,21 @@
 import { GraphQLError } from '../../error/GraphQLError.js';
 
+import type {
+  FragmentDefinitionNode,
+  FragmentSpreadNode,
+  InlineFragmentNode,
+  OperationDefinitionNode,
+  OperationTypeNode,
+  SelectionSetNode,
+} from '../../language/ast.js';
+import { Kind } from '../../language/kinds.js';
 import type { ASTVisitor } from '../../language/visitor.js';
 
 import {
   GraphQLDeferDirective,
   GraphQLStreamDirective,
 } from '../../type/directives.js';
+import type { GraphQLObjectType } from '../../type/index.js';
 
 import type { ValidationContext } from '../ValidationContext.js';
 
@@ -18,46 +28,112 @@ export function DeferStreamDirectiveOnRootFieldRule(
   context: ValidationContext,
 ): ASTVisitor {
   return {
-    Directive(node) {
-      const mutationType = context.getSchema().getMutationType();
-      const subscriptionType = context.getSchema().getSubscriptionType();
-      const parentType = context.getParentType();
-      if (parentType && node.name.value === GraphQLDeferDirective.name) {
-        if (mutationType && parentType === mutationType) {
-          context.reportError(
-            new GraphQLError(
-              `Defer directive cannot be used on root mutation type "${parentType}".`,
-              { nodes: node },
-            ),
-          );
-        }
-        if (subscriptionType && parentType === subscriptionType) {
-          context.reportError(
-            new GraphQLError(
-              `Defer directive cannot be used on root subscription type "${parentType}".`,
-              { nodes: node },
-            ),
-          );
+    OperationDefinition(node: OperationDefinitionNode) {
+      const document = context.getDocument();
+      const fragments = new Map<string, FragmentDefinitionNode>();
+
+      for (const definition of document.definitions) {
+        if (definition.kind === Kind.FRAGMENT_DEFINITION) {
+          fragments.set(definition.name.value, definition);
         }
       }
-      if (parentType && node.name.value === GraphQLStreamDirective.name) {
-        if (mutationType && parentType === mutationType) {
-          context.reportError(
-            new GraphQLError(
-              `Stream directive cannot be used on root mutation type "${parentType}".`,
-              { nodes: node },
-            ),
-          );
-        }
-        if (subscriptionType && parentType === subscriptionType) {
+      if (node.operation !== 'subscription' && node.operation !== 'mutation') {
+        return;
+      }
+      const schema = context.getSchema();
+      const rootType = schema.getRootType(node.operation);
+      if (rootType) {
+        collectRootFields({
+          context,
+          operationType: node.operation,
+          rootType,
+          fragments,
+          selectionSet: node.selectionSet,
+          visitedFragments: new Set(),
+        });
+      }
+    },
+  };
+}
+
+function collectRootFields({
+  context,
+  operationType,
+  rootType,
+  fragments,
+  selectionSet,
+  visitedFragments,
+}: {
+  context: ValidationContext;
+  operationType: OperationTypeNode;
+  rootType: GraphQLObjectType;
+  fragments: Map<string, FragmentDefinitionNode>;
+  selectionSet: SelectionSetNode;
+  visitedFragments: Set<string>;
+}) {
+  for (const selection of selectionSet.selections) {
+    if (selection.kind === 'Field') {
+      const stream = selection.directives?.find(
+        (d) => d.name.value === GraphQLStreamDirective.name,
+      );
+      if (stream) {
+        context.reportError(
+          new GraphQLError(
+            `Stream directive cannot be used on root ${operationType} type "${rootType}".`,
+            { nodes: stream },
+          ),
+        );
+      }
+    } else if (selection.kind === 'FragmentSpread') {
+      const fragmentName = selection.name.value;
+      if (visitedFragments.has(fragmentName)) {
+        continue;
+      }
+      const fragment = fragments.get(fragmentName);
+      if (fragment) {
+        const defer = getDeferDirective(selection);
+        if (defer) {
           context.reportError(
             new GraphQLError(
-              `Stream directive cannot be used on root subscription type "${parentType}".`,
-              { nodes: node },
+              `Defer directive cannot be used on root ${operationType} type "${rootType}".`,
+              { nodes: defer },
             ),
           );
         }
+        collectRootFields({
+          context,
+          operationType,
+          rootType,
+          fragments,
+          selectionSet: fragment.selectionSet,
+          visitedFragments,
+        });
       }
-    },
-  };
+      visitedFragments.add(fragmentName);
+    } else if (selection.kind === 'InlineFragment') {
+      const defer = getDeferDirective(selection);
+      if (defer) {
+        context.reportError(
+          new GraphQLError(
+            `Defer directive cannot be used on root ${operationType} type "${rootType}".`,
+            { nodes: defer },
+          ),
+        );
+      }
+      collectRootFields({
+        context,
+        operationType,
+        rootType,
+        fragments,
+        selectionSet: selection.selectionSet,
+        visitedFragments,
+      });
+    }
+  }
+}
+
+function getDeferDirective(fragment: FragmentSpreadNode | InlineFragmentNode) {
+  return fragment.directives?.find(
+    (d) => d.name.value === GraphQLDeferDirective.name,
+  );
 }