align return types from execution and subscription

with respect to possible promises

Author: yaacovCR

src/execution/__tests__/subscribe-test.ts

@@ -5,13 +5,16 @@ import { expectJSON } from '../../__testUtils__/expectJSON';
 import { resolveOnNextTick } from '../../__testUtils__/resolveOnNextTick';
 
 import { isAsyncIterable } from '../../jsutils/isAsyncIterable';
+import { isPromise } from '../../jsutils/isPromise';
+import type { PromiseOrValue } from '../../jsutils/PromiseOrValue';
 
 import { parse } from '../../language/parser';
 
 import { GraphQLList, GraphQLObjectType } from '../../type/definition';
 import { GraphQLBoolean, GraphQLInt, GraphQLString } from '../../type/scalars';
 import { GraphQLSchema } from '../../type/schema';
 
+import type { ExecutionResult } from '../execute';
 import { createSourceEventStream, subscribe } from '../subscribe';
 
 import { SimplePubSub } from './simplePubSub';
@@ -134,9 +137,6 @@ async function expectPromise(promise: Promise<unknown>) {
   }
 
   return {
-    toReject() {
-      expect(caughtError).to.be.an.instanceOf(Error);
-    },
     toRejectWith(message: string) {
       expect(caughtError).to.be.an.instanceOf(Error);
       expect(caughtError).to.have.property('message', message);
@@ -151,6 +151,34 @@ const DummyQueryType = new GraphQLObjectType({
   },
 });
 
+function subscribeWithFn(
+  subscribeFn: () => unknown,
+): PromiseOrValue<ExecutionResult | AsyncGenerator<ExecutionResult>> {
+  const schema = new GraphQLSchema({
+    query: DummyQueryType,
+    subscription: new GraphQLObjectType({
+      name: 'Subscription',
+      fields: {
+        foo: { type: GraphQLString, subscribe: subscribeFn },
+      },
+    }),
+  });
+  const document = parse('subscription { foo }');
+  const result = subscribe({ schema, document });
+
+  if (isPromise(result)) {
+    return result.then(async (resolvedResult) => {
+      expectJSON(await createSourceEventStream(schema, document)).toDeepEqual(
+        resolvedResult,
+      );
+      return result;
+    });
+  }
+
+  expectJSON(createSourceEventStream(schema, document)).toDeepEqual(result);
+  return result;
+}
+
 /* eslint-disable @typescript-eslint/require-await */
 // Check all error cases when initializing the subscription.
 describe('Subscription Initialization Phase', () => {
@@ -356,24 +384,22 @@ describe('Subscription Initialization Phase', () => {
     });
 
     // @ts-expect-error (schema must not be null)
-    (await expectPromise(subscribe({ schema: null, document }))).toRejectWith(
+    expect(() => subscribe({ schema: null, document })).to.throw(
       'Expected null to be a GraphQL schema.',
     );
 
     // @ts-expect-error
-    (await expectPromise(subscribe({ document }))).toRejectWith(
+    expect(() => subscribe({ document })).to.throw(
       'Expected undefined to be a GraphQL schema.',
     );
 
     // @ts-expect-error (document must not be null)
-    (await expectPromise(subscribe({ schema, document: null }))).toRejectWith(
+    expect(() => subscribe({ schema, document: null })).to.throw(
       'Must provide document.',
     );
 
     // @ts-expect-error
-    (await expectPromise(subscribe({ schema }))).toRejectWith(
-      'Must provide document.',
-    );
+    expect(() => subscribe({ schema })).to.throw('Must provide document.');
   });
 
   it('resolves to an error if schema does not support subscriptions', async () => {
@@ -427,50 +453,22 @@ describe('Subscription Initialization Phase', () => {
     });
 
     // @ts-expect-error
-    (await expectPromise(subscribe({ schema, document: {} }))).toReject();
+    expect(() => subscribe({ schema, document: {} })).to.throw();
   });
 
   it('throws an error if subscribe does not return an iterator', async () => {
-    const schema = new GraphQLSchema({
-      query: DummyQueryType,
-      subscription: new GraphQLObjectType({
-        name: 'Subscription',
-        fields: {
-          foo: {
-            type: GraphQLString,
-            subscribe: () => 'test',
-          },
-        },
-      }),
-    });
-
-    const document = parse('subscription { foo }');
+    expect(() => subscribeWithFn(() => 'test')).to.throw(
+      'Subscription field must return Async Iterable. Received: "test".',
+    );
 
-    (await expectPromise(subscribe({ schema, document }))).toRejectWith(
+    const result = subscribeWithFn(() => Promise.resolve('test'));
+    assert(isPromise(result));
+    (await expectPromise(result)).toRejectWith(
       'Subscription field must return Async Iterable. Received: "test".',
     );
   });
 
   it('resolves to an error for subscription resolver errors', async () => {
-    async function subscribeWithFn(subscribeFn: () => unknown) {
-      const schema = new GraphQLSchema({
-        query: DummyQueryType,
-        subscription: new GraphQLObjectType({
-          name: 'Subscription',
-          fields: {
-            foo: { type: GraphQLString, subscribe: subscribeFn },
-          },
-        }),
-      });
-      const document = parse('subscription { foo }');
-      const result = await subscribe({ schema, document });
-
-      expectJSON(await createSourceEventStream(schema, document)).toDeepEqual(
-        result,
-      );
-      return result;
-    }
-
     const expectedResult = {
       errors: [
         {
@@ -483,12 +481,12 @@ describe('Subscription Initialization Phase', () => {
 
     expectJSON(
       // Returning an error
-      await subscribeWithFn(() => new Error('test error')),
+      subscribeWithFn(() => new Error('test error')),
     ).toDeepEqual(expectedResult);
 
     expectJSON(
       // Throwing an error
-      await subscribeWithFn(() => {
+      subscribeWithFn(() => {
         throw new Error('test error');
       }),
     ).toDeepEqual(expectedResult);

src/execution/subscribe.ts

@@ -1,7 +1,9 @@
 import { inspect } from '../jsutils/inspect';
 import { isAsyncIterable } from '../jsutils/isAsyncIterable';
+import { isPromise } from '../jsutils/isPromise';
 import type { Maybe } from '../jsutils/Maybe';
 import { addPath, pathToArray } from '../jsutils/Path';
+import type { PromiseOrValue } from '../jsutils/PromiseOrValue';
 
 import { GraphQLError } from '../error/GraphQLError';
 import { locatedError } from '../error/locatedError';
@@ -47,9 +49,11 @@ import { getArgumentValues } from './values';
  *
  * Accepts either an object with named arguments, or individual arguments.
  */
-export async function subscribe(
+export function subscribe(
   args: ExecutionArgs,
-): Promise<AsyncGenerator<ExecutionResult, void, void> | ExecutionResult> {
+): PromiseOrValue<
+  AsyncGenerator<ExecutionResult, void, void> | ExecutionResult
+> {
   const {
     schema,
     document,
@@ -61,7 +65,7 @@ export async function subscribe(
     subscribeFieldResolver,
   } = args;
 
-  const resultOrStream = await createSourceEventStream(
+  const resultOrStream = createSourceEventStream(
     schema,
     document,
     rootValue,
@@ -71,6 +75,42 @@ export async function subscribe(
     subscribeFieldResolver,
   );
 
+  if (isPromise(resultOrStream)) {
+    return resultOrStream.then((resolvedResultOrStream) =>
+      mapSourceToResponse(
+        schema,
+        document,
+        resolvedResultOrStream,
+        contextValue,
+        variableValues,
+        operationName,
+        fieldResolver,
+      ),
+    );
+  }
+
+  return mapSourceToResponse(
+    schema,
+    document,
+    resultOrStream,
+    contextValue,
+    variableValues,
+    operationName,
+    fieldResolver,
+  );
+}
+
+function mapSourceToResponse(
+  schema: GraphQLSchema,
+  document: DocumentNode,
+  resultOrStream: ExecutionResult | AsyncIterable<unknown>,
+  contextValue?: unknown,
+  variableValues?: Maybe<{ readonly [variable: string]: unknown }>,
+  operationName?: Maybe<string>,
+  fieldResolver?: Maybe<GraphQLFieldResolver<any, any>>,
+): PromiseOrValue<
+  AsyncGenerator<ExecutionResult, void, void> | ExecutionResult
+> {
   if (!isAsyncIterable(resultOrStream)) {
     return resultOrStream;
   }
@@ -81,7 +121,7 @@ export async function subscribe(
   // 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.
-  const mapSourceToResponse = (payload: unknown) =>
+  return mapAsyncIterator(resultOrStream, (payload: unknown) =>
     execute({
       schema,
       document,
@@ -90,10 +130,8 @@ export async function subscribe(
       variableValues,
       operationName,
       fieldResolver,
-    });
-
-  // Map every source value to a ExecutionResult value as described above.
-  return mapAsyncIterator(resultOrStream, mapSourceToResponse);
+    }),
+  );
 }
 
 /**
@@ -124,15 +162,15 @@ export async function subscribe(
  * or otherwise separating these two steps. For more on this, see the
  * "Supporting Subscriptions at Scale" information in the GraphQL specification.
  */
-export async function createSourceEventStream(
+export function createSourceEventStream(
   schema: GraphQLSchema,
   document: DocumentNode,
   rootValue?: unknown,
   contextValue?: unknown,
   variableValues?: Maybe<{ readonly [variable: string]: unknown }>,
   operationName?: Maybe<string>,
   subscribeFieldResolver?: Maybe<GraphQLFieldResolver<any, any>>,
-): Promise<AsyncIterable<unknown> | ExecutionResult> {
+): PromiseOrValue<AsyncIterable<unknown> | ExecutionResult> {
   // If arguments are missing or incorrectly typed, this is an internal
   // developer mistake which should throw an early error.
   assertValidExecutionArguments(schema, document, variableValues);
@@ -155,17 +193,22 @@ export async function createSourceEventStream(
   }
 
   try {
-    const eventStream = await executeSubscription(exeContext);
-
-    // Assert field returned an event stream, otherwise yield an error.
-    if (!isAsyncIterable(eventStream)) {
-      throw new Error(
-        'Subscription field must return Async Iterable. ' +
-          `Received: ${inspect(eventStream)}.`,
-      );
+    const eventStream = executeSubscription(exeContext);
+
+    if (isPromise(eventStream)) {
+      return eventStream
+        .then((resolvedEventStream) => ensureAsyncIterable(resolvedEventStream))
+        .then(undefined, (error) => {
+          // If it GraphQLError, report it as an ExecutionResult, containing only errors and no data.
+          // Otherwise treat the error as a system-class error and re-throw it.
+          if (error instanceof GraphQLError) {
+            return { errors: [error] };
+          }
+          throw error;
+        });
     }
 
-    return eventStream;
+    return ensureAsyncIterable(eventStream);
   } catch (error) {
     // If it GraphQLError, report it as an ExecutionResult, containing only errors and no data.
     // Otherwise treat the error as a system-class error and re-throw it.
@@ -176,9 +219,19 @@ export async function createSourceEventStream(
   }
 }
 
-async function executeSubscription(
-  exeContext: ExecutionContext,
-): Promise<unknown> {
+function ensureAsyncIterable(eventStream: unknown): AsyncIterable<unknown> {
+  // Assert field returned an event stream, otherwise yield an error.
+  if (!isAsyncIterable(eventStream)) {
+    throw new Error(
+      'Subscription field must return Async Iterable. ' +
+        `Received: ${inspect(eventStream)}.`,
+    );
+  }
+
+  return eventStream;
+}
+
+function executeSubscription(exeContext: ExecutionContext): unknown {
   const { schema, fragments, operation, variableValues, rootValue } =
     exeContext;
 
@@ -233,13 +286,26 @@ async function executeSubscription(
     // Call the `subscribe()` resolver or the default resolver to produce an
     // AsyncIterable yielding raw payloads.
     const resolveFn = fieldDef.subscribe ?? exeContext.subscribeFieldResolver;
-    const eventStream = await resolveFn(rootValue, args, contextValue, info);
 
-    if (eventStream instanceof Error) {
-      throw eventStream;
+    const eventStream = resolveFn(rootValue, args, contextValue, info);
+
+    if (isPromise(eventStream)) {
+      return eventStream
+        .then((resolvedEventStream) => throwReturnedError(resolvedEventStream))
+        .then(undefined, (error) => {
+          throw locatedError(error, fieldNodes, pathToArray(path));
+        });
     }
-    return eventStream;
+
+    return throwReturnedError(eventStream);
   } catch (error) {
     throw locatedError(error, fieldNodes, pathToArray(path));
   }
 }
+
+function throwReturnedError(eventStream: unknown): unknown {
+  if (eventStream instanceof Error) {
+    throw eventStream;
+  }
+  return eventStream;
+}