Raise a field error when encountering defer/stream directive while executing subscription

Author: robrichard

src/execution/__tests__/defer-test.ts

@@ -670,7 +670,7 @@ describe('Execute: defer directive', () => {
         rootValue: {},
       }),
     ).to.throw(
-      'Executing this GraphQL operation would unexpectedly produce multiple payloads (due to @defer or @stream directive). Disable `@defer` or `@stream` by setting the `if` argument to `false`.',
+      'Executing this GraphQL operation would unexpectedly produce multiple payloads (due to @defer or @stream directive)',
     );
   });
 
@@ -693,7 +693,7 @@ describe('Execute: defer directive', () => {
       errors: [
         {
           message:
-            'Executing this GraphQL operation would unexpectedly produce multiple payloads (due to @defer or @stream directive). Disable `@defer` or `@stream` by setting the `if` argument to `false`.',
+            'Executing this GraphQL operation would unexpectedly produce multiple payloads (due to @defer or @stream directive)',
         },
       ],
     });

src/execution/__tests__/subscribe-test.ts

@@ -97,7 +97,12 @@ function createSubscription(
   variableValues?: { readonly [variable: string]: unknown },
 ) {
   const document = parse(`
-    subscription ($priority: Int = 0, $shouldDefer: Boolean = false, $asyncResolver: Boolean = false) {
+    subscription (
+      $priority: Int = 0
+      $shouldDefer: Boolean = false
+      $shouldStream: Boolean = false
+      $asyncResolver: Boolean = false
+    ) {
       importantEmail(priority: $priority) {
         email {
           from
@@ -108,6 +113,7 @@ function createSubscription(
         }
         ... @defer(if: $shouldDefer) {
           inbox {
+            emails @include(if: $shouldStream) @stream(if: $shouldStream)
             unread
             total
           }
@@ -723,9 +729,12 @@ describe('Subscription Publish Phase', () => {
         errors: [
           {
             message:
-              'Executing this GraphQL operation would unexpectedly produce multiple payloads (due to @defer or @stream directive). Disable `@defer` or `@stream` by setting the `if` argument to `false`.',
+              '`@defer` directive not supported on subscription operations. Disable `@defer` by setting the `if` argument to `false`.',
+            locations: [{ line: 3, column: 7 }],
+            path: ['importantEmail'],
           },
         ],
+        data: { importantEmail: null },
       },
     };
 
@@ -757,6 +766,89 @@ describe('Subscription Publish Phase', () => {
     });
   });
 
+  it('subscribe function returns errors with @stream', async () => {
+    const pubsub = new SimplePubSub<Email>();
+    const subscription = await createSubscription(pubsub, {
+      shouldStream: true,
+    });
+    assert(isAsyncIterable(subscription));
+    // Wait for the next subscription payload.
+    const payload = subscription.next();
+
+    // A new email arrives!
+    expect(
+      pubsub.emit({
+        from: '[email protected]',
+        subject: 'Alright',
+        message: 'Tests are good',
+        unread: true,
+      }),
+    ).to.equal(true);
+
+    // The previously waited on payload now has a value.
+    expectJSON(await payload).toDeepEqual({
+      done: false,
+      value: {
+        errors: [
+          {
+            message:
+              '`@stream` directive not supported on subscription operations. Disable `@stream` by setting the `if` argument to `false`.',
+            locations: [{ line: 13, column: 13 }],
+            path: ['importantEmail', 'inbox', 'emails'],
+          },
+        ],
+        data: {
+          importantEmail: {
+            email: { from: '[email protected]', subject: 'Alright' },
+            inbox: { emails: null, unread: 1, total: 2 },
+          },
+        },
+      },
+    });
+
+    // Another new email arrives, after all incrementally delivered payloads are received.
+    expect(
+      pubsub.emit({
+        from: '[email protected]',
+        subject: 'Tools',
+        message: 'I <3 making things',
+        unread: true,
+      }),
+    ).to.equal(true);
+
+    // The next waited on payload will have a value.
+    expectJSON(await subscription.next()).toDeepEqual({
+      done: false,
+      value: {
+        errors: [
+          {
+            message:
+              '`@stream` directive not supported on subscription operations. Disable `@stream` by setting the `if` argument to `false`.',
+            locations: [{ line: 13, column: 13 }],
+            path: ['importantEmail', 'inbox', 'emails'],
+          },
+        ],
+        data: {
+          importantEmail: {
+            email: { from: '[email protected]', subject: 'Tools' },
+            inbox: { emails: null, unread: 2, total: 3 },
+          },
+        },
+      },
+    });
+
+    expectJSON(await subscription.return()).toDeepEqual({
+      done: true,
+      value: undefined,
+    });
+
+    // Awaiting a subscription after closing it results in completed results.
+    expectJSON(await subscription.next()).toDeepEqual({
+      done: true,
+      value: undefined,
+    });
+  });
+
   it('produces a payload when there are multiple events', async () => {
     const pubsub = new SimplePubSub<Email>();
     const subscription = createSubscription(pubsub);

src/execution/collectFields.ts

@@ -1,13 +1,16 @@
 import { AccumulatorMap } from '../jsutils/AccumulatorMap.js';
+import { invariant } from '../jsutils/invariant.js';
 import type { ObjMap } from '../jsutils/ObjMap.js';
 
 import type {
   FieldNode,
   FragmentDefinitionNode,
   FragmentSpreadNode,
   InlineFragmentNode,
+  OperationDefinitionNode,
   SelectionSetNode,
 } from '../language/ast.js';
+import { OperationTypeNode } from '../language/ast.js';
 import { Kind } from '../language/kinds.js';
 
 import type { GraphQLObjectType } from '../type/definition.js';
@@ -47,16 +50,17 @@ export function collectFields(
   fragments: ObjMap<FragmentDefinitionNode>,
   variableValues: { [variable: string]: unknown },
   runtimeType: GraphQLObjectType,
-  selectionSet: SelectionSetNode,
+  operation: OperationDefinitionNode,
 ): FieldsAndPatches {
   const fields = new AccumulatorMap<string, FieldNode>();
   const patches: Array<PatchFields> = [];
   collectFieldsImpl(
     schema,
     fragments,
     variableValues,
+    operation,
     runtimeType,
-    selectionSet,
+    operation.selectionSet,
     fields,
     patches,
     new Set(),
@@ -74,10 +78,12 @@ export function collectFields(
  *
  * @internal
  */
+// eslint-disable-next-line max-params
 export function collectSubfields(
   schema: GraphQLSchema,
   fragments: ObjMap<FragmentDefinitionNode>,
   variableValues: { [variable: string]: unknown },
+  operation: OperationDefinitionNode,
   returnType: GraphQLObjectType,
   fieldNodes: ReadonlyArray<FieldNode>,
 ): FieldsAndPatches {
@@ -96,6 +102,7 @@ export function collectSubfields(
         schema,
         fragments,
         variableValues,
+        operation,
         returnType,
         node.selectionSet,
         subFieldNodes,
@@ -112,6 +119,7 @@ function collectFieldsImpl(
   schema: GraphQLSchema,
   fragments: ObjMap<FragmentDefinitionNode>,
   variableValues: { [variable: string]: unknown },
+  operation: OperationDefinitionNode,
   runtimeType: GraphQLObjectType,
   selectionSet: SelectionSetNode,
   fields: AccumulatorMap<string, FieldNode>,
@@ -135,14 +143,15 @@ function collectFieldsImpl(
           continue;
         }
 
-        const defer = getDeferValues(variableValues, selection);
+        const defer = getDeferValues(operation, variableValues, selection);
 
         if (defer) {
           const patchFields = new AccumulatorMap<string, FieldNode>();
           collectFieldsImpl(
             schema,
             fragments,
             variableValues,
+            operation,
             runtimeType,
             selection.selectionSet,
             patchFields,
@@ -158,6 +167,7 @@ function collectFieldsImpl(
             schema,
             fragments,
             variableValues,
+            operation,
             runtimeType,
             selection.selectionSet,
             fields,
@@ -174,7 +184,7 @@ function collectFieldsImpl(
           continue;
         }
 
-        const defer = getDeferValues(variableValues, selection);
+        const defer = getDeferValues(operation, variableValues, selection);
         if (visitedFragmentNames.has(fragName) && !defer) {
           continue;
         }
@@ -197,6 +207,7 @@ function collectFieldsImpl(
             schema,
             fragments,
             variableValues,
+            operation,
             runtimeType,
             fragment.selectionSet,
             patchFields,
@@ -212,6 +223,7 @@ function collectFieldsImpl(
             schema,
             fragments,
             variableValues,
+            operation,
             runtimeType,
             fragment.selectionSet,
             fields,
@@ -231,6 +243,7 @@ function collectFieldsImpl(
  * not disabled by the "if" argument.
  */
 function getDeferValues(
+  operation: OperationDefinitionNode,
   variableValues: { [variable: string]: unknown },
   node: FragmentSpreadNode | InlineFragmentNode,
 ): undefined | { label: string | undefined } {
@@ -244,6 +257,11 @@ function getDeferValues(
     return;
   }
 
+  invariant(
+    operation.operation !== OperationTypeNode.SUBSCRIPTION,
+    '`@defer` directive not supported on subscription operations. Disable `@defer` by setting the `if` argument to `false`.',
+  );
+
   return {
     label: typeof defer.label === 'string' ? defer.label : undefined,
   };

src/execution/execute.ts

@@ -78,6 +78,7 @@ const collectSubfields = memoize3(
       exeContext.schema,
       exeContext.fragments,
       exeContext.variableValues,
+      exeContext.operation,
       returnType,
       fieldNodes,
     ),
@@ -263,7 +264,7 @@ export interface ExecutionArgs {
 }
 
 const UNEXPECTED_MULTIPLE_PAYLOADS =
-  'Executing this GraphQL operation would unexpectedly produce multiple payloads (due to @defer or @stream directive). Disable `@defer` or `@stream` by setting the `if` argument to `false`.';
+  'Executing this GraphQL operation would unexpectedly produce multiple payloads (due to @defer or @stream directive)';
 
 /**
  * Implements the "Executing requests" section of the GraphQL specification.
@@ -530,7 +531,7 @@ function executeOperation(
     fragments,
     variableValues,
     rootType,
-    operation.selectionSet,
+    operation,
   );
   const path = undefined;
   let result;
@@ -949,6 +950,11 @@ function getStreamValues(
     'initialCount must be a positive integer',
   );
 
+  invariant(
+    exeContext.operation.operation !== OperationTypeNode.SUBSCRIPTION,
+    '`@stream` directive not supported on subscription operations. Disable `@stream` by setting the `if` argument to `false`.',
+  );
+
   return {
     initialCount: stream.initialCount,
     label: typeof stream.label === 'string' ? stream.label : undefined,
@@ -1456,17 +1462,6 @@ export const defaultFieldResolver: GraphQLFieldResolver<unknown, unknown> =
     }
   };
 
-function ensureSingleExecutionResult(
-  result: ExecutionResult | ExperimentalIncrementalExecutionResults,
-): ExecutionResult {
-  if ('initialResult' in result) {
-    return {
-      errors: [new GraphQLError(UNEXPECTED_MULTIPLE_PAYLOADS)],
-    };
-  }
-  return result;
-}
-
 /**
  * Implements the "Subscribe" algorithm described in the GraphQL specification.
  *
@@ -1534,10 +1529,15 @@ function mapSourceToResponse(
   // the GraphQL specification. The `execute` function provides the
   // "ExecuteSubscriptionEvent" algorithm, as it is nearly identical to the
   // "ExecuteQuery" algorithm, for which `execute` is also used.
-  return mapAsyncIterable(resultOrStream, async (payload: unknown) =>
-    ensureSingleExecutionResult(
-      await executeImpl(buildPerEventExecutionContext(exeContext, payload)),
-    ),
+  return mapAsyncIterable(
+    resultOrStream,
+    (payload: unknown) =>
+      executeImpl(
+        buildPerEventExecutionContext(exeContext, payload),
+        // typecast to ExecutionResult, not possible to return
+        // ExperimentalIncrementalExecutionResults when
+        // exeContext.operation is 'subscription'.
+      ) as ExecutionResult,
   );
 }
 
@@ -1618,7 +1618,7 @@ function executeSubscription(
     fragments,
     variableValues,
     rootType,
-    operation.selectionSet,
+    operation,
   );
 
   const firstRootField = rootFields.entries().next().value;

src/validation/rules/SingleFieldSubscriptionsRule.ts

@@ -46,7 +46,7 @@ export function SingleFieldSubscriptionsRule(
             fragments,
             variableValues,
             subscriptionType,
-            node.selectionSet,
+            node,
           );
           if (fields.size > 1) {
             const fieldSelectionLists = [...fields.values()];