Implement support for @stream directive

Author: robrichard

src/execution/__tests__/stream-test.js

@@ -50,6 +50,7 @@ import {
   GraphQLIncludeDirective,
   GraphQLSkipDirective,
   GraphQLDeferDirective,
+  GraphQLStreamDirective,
 } from '../type/directives';
 import {
   isObjectType,
@@ -701,6 +702,42 @@ function getDeferValues(
   };
 }
 
+/**
+ * Returns an object containing the @stream arguments if a field should be
+ * streamed based on the experimental flag, stream directive present and
+ * not disabled by the "if" argument.
+ */
+function getStreamValues(
+  exeContext: ExecutionContext,
+  fieldNodes: $ReadOnlyArray<FieldNode>,
+): void | {|
+  initialCount?: number,
+  label?: string,
+|} {
+  // validation only allows equivalent streams on multiple fields, so it is
+  // safe to only check the first fieldNode for the stream directive
+  const stream = getDirectiveValues(
+    GraphQLStreamDirective,
+    fieldNodes[0],
+    exeContext.variableValues,
+  );
+
+  if (!stream) {
+    return;
+  }
+
+  if (stream.if === false) {
+    return;
+  }
+
+  return {
+    initialCount:
+      // istanbul ignore next (initialCount is required number argument)
+      typeof stream.initialCount === 'number' ? stream.initialCount : undefined,
+    label: typeof stream.label === 'string' ? stream.label : undefined,
+  };
+}
+
 /**
  * Determines if a fragment is applicable to the given type.
  */
@@ -993,6 +1030,7 @@ function completeAsyncIteratorValue(
   errors: Array<GraphQLError>,
 ): Promise<$ReadOnlyArray<mixed>> {
   let containsPromise = false;
+  const stream = getStreamValues(exeContext, fieldNodes);
   return new Promise((resolve) => {
     function next(index, completedResults) {
       const fieldPath = addPath(path, index, undefined);
@@ -1029,7 +1067,26 @@ function completeAsyncIteratorValue(
             return;
           }
 
-          next(index + 1, completedResults);
+          const newIndex = index + 1;
+          if (
+            stream &&
+            typeof stream.initialCount === 'number' &&
+            newIndex >= stream.initialCount
+          ) {
+            exeContext.dispatcher.addAsyncIteratorValue(
+              stream.label,
+              newIndex,
+              path,
+              iterator,
+              exeContext,
+              fieldNodes,
+              info,
+              itemType,
+            );
+            resolve(completedResults);
+            return;
+          }
+          next(newIndex, completedResults);
         },
         (rawError) => {
           completedResults.push(null);
@@ -1084,15 +1141,37 @@ function completeListValue(
     );
   }
 
+  const stream = getStreamValues(exeContext, fieldNodes);
+
   // This is specified as a simple map, however we're optimizing the path
   // where the list contains no Promises by avoiding creating another Promise.
   let containsPromise = false;
-  const completedResults = Array.from(result, (item, index) => {
+  const completedResults = [];
+  let index = 0;
+  for (const item of result) {
     // No need to modify the info object containing the path,
     // since from here on it is not ever accessed by resolver functions.
     const itemPath = addPath(path, index, undefined);
     try {
       let completedItem;
+
+      if (
+        stream &&
+        typeof stream.initialCount === 'number' &&
+        index >= stream.initialCount
+      ) {
+        exeContext.dispatcher.addValue(
+          stream.label,
+          itemPath,
+          item,
+          exeContext,
+          fieldNodes,
+          info,
+          itemType,
+        );
+        index++;
+        continue;
+      }
       if (isPromise(item)) {
         completedItem = item.then((resolved) =>
           completeValue(
@@ -1121,21 +1200,25 @@ function completeListValue(
         containsPromise = true;
         // Note: we don't rely on a `catch` method, but we do expect "thenable"
         // to take a second callback for the error case.
-        return completedItem.then(undefined, (rawError) => {
-          const error = locatedError(
-            rawError,
-            fieldNodes,
-            pathToArray(itemPath),
-          );
-          return handleFieldError(error, itemType, errors);
-        });
+        completedResults.push(
+          completedItem.then(undefined, (rawError) => {
+            const error = locatedError(
+              rawError,
+              fieldNodes,
+              pathToArray(itemPath),
+            );
+            return handleFieldError(error, itemType, errors);
+          }),
+        );
+      } else {
+        completedResults.push(completedItem);
       }
-      return completedItem;
     } catch (rawError) {
       const error = locatedError(rawError, fieldNodes, pathToArray(itemPath));
-      return handleFieldError(error, itemType, errors);
+      completedResults.push(handleFieldError(error, itemType, errors));
     }
-  });
+    index++;
+  }
 
   return containsPromise ? Promise.all(completedResults) : completedResults;
 }
@@ -1551,6 +1634,129 @@ export class Dispatcher {
     );
   }
 
+  addValue(
+    label?: string,
+    path: Path,
+    promiseOrData: PromiseOrValue<ObjMap<mixed> | mixed>,
+    exeContext: ExecutionContext,
+    fieldNodes: $ReadOnlyArray<FieldNode>,
+    info: GraphQLResolveInfo,
+    itemType: GraphQLOutputType,
+  ): void {
+    const errors = [];
+    this._subsequentPayloads.push(
+      Promise.resolve(promiseOrData)
+        .then((resolved) =>
+          completeValue(
+            exeContext,
+            itemType,
+            fieldNodes,
+            info,
+            path,
+            resolved,
+            errors,
+          ),
+        )
+        // Note: we don't rely on a `catch` method, but we do expect "thenable"
+        // to take a second callback for the error case.
+        .then(undefined, (rawError) => {
+          const error = locatedError(rawError, fieldNodes, pathToArray(path));
+          return handleFieldError(error, itemType, errors);
+        })
+        .then((data) => ({
+          value: createPatchResult(data, label, path, errors),
+          done: false,
+        })),
+    );
+  }
+
+  addAsyncIteratorValue(
+    label?: string,
+    initialIndex: number,
+    path?: Path,
+    iterator: AsyncIterator<mixed>,
+    exeContext: ExecutionContext,
+    fieldNodes: $ReadOnlyArray<FieldNode>,
+    info: GraphQLResolveInfo,
+    itemType: GraphQLOutputType,
+  ): void {
+    const subsequentPayloads = this._subsequentPayloads;
+    function next(index) {
+      const fieldPath = addPath(path, index);
+      const patchErrors = [];
+      subsequentPayloads.push(
+        iterator.next().then(
+          ({ value: data, done }) => {
+            if (done) {
+              return { value: undefined, done: true };
+            }
+
+            // eslint-disable-next-line node/callback-return
+            next(index + 1);
+
+            try {
+              const completedItem = completeValue(
+                exeContext,
+                itemType,
+                fieldNodes,
+                info,
+                fieldPath,
+                data,
+                patchErrors,
+              );
+
+              if (isPromise(completedItem)) {
+                return completedItem.then((resolveItem) => ({
+                  value: createPatchResult(
+                    resolveItem,
+                    label,
+                    fieldPath,
+                    patchErrors,
+                  ),
+                  done: false,
+                }));
+              }
+
+              return {
+                value: createPatchResult(
+                  completedItem,
+                  label,
+                  fieldPath,
+                  patchErrors,
+                ),
+                done: false,
+              };
+            } catch (rawError) {
+              const error = locatedError(
+                rawError,
+                fieldNodes,
+                pathToArray(fieldPath),
+              );
+              handleFieldError(error, itemType, patchErrors);
+              return {
+                value: createPatchResult(null, label, fieldPath, patchErrors),
+                done: false,
+              };
+            }
+          },
+          (rawError) => {
+            const error = locatedError(
+              rawError,
+              fieldNodes,
+              pathToArray(fieldPath),
+            );
+            handleFieldError(error, itemType, patchErrors);
+            return {
+              value: createPatchResult(null, label, fieldPath, patchErrors),
+              done: false,
+            };
+          },
+        ),
+      );
+    }
+    next(initialIndex);
+  }
+
   _race(): Promise<IteratorResult<ExecutionPatchResult, void>> {
     return new Promise((resolve) => {
       this._subsequentPayloads.forEach((promise) => {
@@ -1567,7 +1773,20 @@ export class Dispatcher {
         );
         return promise;
       })
-      .then(({ value }) => {
+      .then(({ value, done }) => {
+        if (done && this._subsequentPayloads.length === 0) {
+          // async iterable resolver just finished and no more pending payloads
+          return {
+            value: {
+              hasNext: false,
+            },
+            done: false,
+          };
+        } else if (done) {
+          // async iterable resolver just finished but there are pending payloads
+          // return the next one
+          return this._race();
+        }
         const returnValue: ExecutionPatchResult = {
           ...value,
           hasNext: this._subsequentPayloads.length > 0,