feat(server): Pass roots for operation fields as an option

src/server.ts

@@ -7,6 +7,7 @@
 import * as http from 'http';
 import * as WebSocket from 'ws';
 import {
+  OperationTypeNode,
   GraphQLSchema,
   ValidationRule,
   ExecutionResult,
@@ -51,6 +52,16 @@ export interface ServerOptions {
    * from the `onSubscribe` callback.
    */
   schema?: GraphQLSchema;
+  /**
+   * The GraphQL root fields or resolvers to go
+   * alongside the schema. Learn more about them
+   * here: https://graphql.org/learn/execution/#root-fields-resolvers.
+   * Related operation root value will be injected to the
+   * `ExecutionArgs` BEFORE the `onSubscribe` callback.
+   */
+  roots?: {
+    [operation in OperationTypeNode]?: Record<string, unknown>;
+  };
   /**
    * Is the `subscribe` function
    * from GraphQL which is used to
@@ -199,6 +210,7 @@ export function createServer(
 ): Server {
   const {
     schema,
+    roots,
     execute,
     onConnect,
     connectionInitWaitTimeout = 3 * 1000, // 3 seconds
@@ -371,17 +383,30 @@ export function createServer(
             }
 
             const operation = message.payload;
+            const document =
+              typeof operation.query === 'string'
+                ? parse(operation.query)
+                : operation.query;
+            const operationAST = getOperationAST(
+              document,
+              operation.operationName,
+            );
+            if (!operationAST) {
+              throw new Error('Unable to get operation AST');
+            }
 
             let execArgsMaybeSchema: Optional<ExecutionArgs, 'schema'> = {
               schema,
               operationName: operation.operationName,
-              document:
-                typeof operation.query === 'string'
-                  ? parse(operation.query)
-                  : operation.query,
+              document,
               variableValues: operation.variables,
             };
 
+            // if roots are provided, inject the coresponding operation root
+            if (roots) {
+              execArgsMaybeSchema.rootValue = roots[operationAST.operation];
+            }
+
             let onSubscribeFormatter: ExecutionResultFormatter | undefined;
             if (onSubscribe) {
               [execArgsMaybeSchema, onSubscribeFormatter] = await onSubscribe(
@@ -415,14 +440,7 @@ export function createServer(
               });
             }
 
-            // execute
-            const operationAST = getOperationAST(
-              execArgs.document,
-              execArgs.operationName,
-            );
-            if (!operationAST) {
-              throw new Error('Unable to get operation AST');
-            }
+            // perform
             if (operationAST.operation === 'subscription') {
               const subscriptionOrResult = await subscribe(execArgs);
               if (isAsyncIterable(subscriptionOrResult)) {

src/tests/server.ts

@@ -1,5 +1,5 @@
 import WebSocket from 'ws';
-import { parse } from 'graphql';
+import { execute, parse, buildSchema } from 'graphql';
 import { GRAPHQL_TRANSPORT_WS_PROTOCOL } from '../protocol';
 import { MessageType, parseMessage, stringifyMessage } from '../message';
 import { startServer, url, schema, pubsub } from './fixtures/simple';
@@ -757,3 +757,97 @@ describe('keepAlive', () => {
     expect(client.readyState).toBe(WebSocket.CLOSED);
   });
 });
+
+it('should use the provided roots as resolvers', async () => {
+  const schema = buildSchema(`
+    type Query {
+      hello: String
+    }
+    type Subscription {
+      count: Int
+    }
+  `);
+
+  const roots = {
+    query: {
+      hello: () => 'Hello World!',
+    },
+    subscription: {
+      count: async function* () {
+        for (const num of [1, 2, 3]) {
+          yield num;
+        }
+      },
+    },
+  };
+
+  await makeServer({
+    schema,
+    roots,
+  });
+
+  const client = new WebSocket(url, GRAPHQL_TRANSPORT_WS_PROTOCOL);
+  client.onopen = () => {
+    client.send(
+      stringifyMessage<MessageType.ConnectionInit>({
+        type: MessageType.ConnectionInit,
+      }),
+    );
+  };
+  await wait(10);
+
+  await new Promise((resolve, reject) => {
+    client.send(
+      stringifyMessage<MessageType.Subscribe>({
+        id: '1',
+        type: MessageType.Subscribe,
+        payload: {
+          query: `{ hello }`,
+        },
+      }),
+    );
+    client.on('message', function onMessage(data) {
+      const message = parseMessage(data);
+      switch (message.type) {
+        case MessageType.Next:
+          expect(message.type).toBe(MessageType.Next);
+          expect(message.payload).toEqual({ data: { hello: 'Hello World!' } });
+          break;
+        case MessageType.Error:
+          client.off('message', onMessage);
+          return reject();
+        case MessageType.Complete:
+          client.off('message', onMessage);
+          return resolve();
+      }
+    });
+  });
+
+  const nextFn = jest.fn();
+  await new Promise((resolve, reject) => {
+    client.send(
+      stringifyMessage<MessageType.Subscribe>({
+        id: '2',
+        type: MessageType.Subscribe,
+        payload: {
+          query: `subscription { count }`,
+        },
+      }),
+    );
+    client.on('message', function onMessage(data) {
+      const message = parseMessage(data);
+      switch (message.type) {
+        case MessageType.Next:
+          nextFn();
+          break;
+        case MessageType.Error:
+          client.off('message', onMessage);
+          return reject(message.payload);
+        case MessageType.Complete:
+          client.off('message', onMessage);
+          return resolve();
+      }
+    });
+  });
+  expect(nextFn).toBeCalledTimes(3);
+});