feat(server): Dynamic schema support by accepting a function or a Promise

README.md

@@ -1009,6 +1009,41 @@ useServer(
 
 </details>
 
+<details id="dynamic-schema">
+<summary><a href="#dynamic-schema">🔗</a> <a href="https://github.com/websockets/ws">ws</a> server usage with dynamic schema</summary>
+
+```typescript
+import { execute, subscribe } from 'graphql';
+import ws from 'ws'; // yarn add ws
+import { useServer } from 'graphql-ws/lib/use/ws';
+import { schema, checkIsAdmin, getDebugSchema } from './my-graphql';
+
+const wsServer = new ws.Server({
+  port: 443,
+  path: '/graphql',
+});
+
+useServer(
+  {
+    execute,
+    subscribe,
+    schema: async (ctx, msg, executionArgsWithoutSchema) => {
+      // will be called on every subscribe request
+      // allowing you to dynamically supply the schema
+      // using the depending on the provided arguments.
+      // throwing an error here closes the socket with
+      // the `Error` message in the close event reason
+      const isAdmin = await checkIsAdmin(ctx.request);
+      if (isAdmin) return getDebugSchema(ctx, msg, executionArgsWithoutSchema);
+      return schema;
+    },
+  },
+  wsServer,
+);
+```
+
+</details>
+
 <details id="custom-exec">
 <summary><a href="#custom-exec">🔗</a> <a href="https://github.com/websockets/ws">ws</a> server usage with custom execution arguments and validation</summary>
 

docs/interfaces/server.serveroptions.md

@@ -410,15 +410,23 @@ ___
 
 ### schema
 
-• `Optional` **schema**: *GraphQLSchema*
+• `Optional` **schema**: *GraphQLSchema* \| (`ctx`: [*Context*](server.context.md)<E\>, `message`: [*SubscribeMessage*](message.subscribemessage.md), `args`: *Omit*<ExecutionArgs, *schema*\>) => *GraphQLSchema* \| *Promise*<GraphQLSchema\>
 
 The GraphQL schema on which the operations
 will be executed and validated against.
 
+If a function is provided, it will be called on
+every subscription request allowing you to manipulate
+schema dynamically.
+
 If the schema is left undefined, you're trusted to
 provide one in the returned `ExecutionArgs` from the
 `onSubscribe` callback.
 
+Throwing an error from within this function will
+close the socket with the `Error` message
+in the close event reason.
+
 ___
 
 ### subscribe

src/server.ts

@@ -57,11 +57,25 @@ export interface ServerOptions<E = unknown> {
    * The GraphQL schema on which the operations
    * will be executed and validated against.
    *
+   * If a function is provided, it will be called on
+   * every subscription request allowing you to manipulate
+   * schema dynamically.
+   *
    * If the schema is left undefined, you're trusted to
    * provide one in the returned `ExecutionArgs` from the
    * `onSubscribe` callback.
+   *
+   * Throwing an error from within this function will
+   * close the socket with the `Error` message
+   * in the close event reason.
    */
-  schema?: GraphQLSchema;
+  schema?:
+    | GraphQLSchema
+    | ((
+        ctx: Context<E>,
+        message: SubscribeMessage,
+        args: Omit<ExecutionArgs, 'schema'>,
+      ) => Promise<GraphQLSchema> | GraphQLSchema);
   /**
    * A value which is provided to every resolver and holds
    * important contextual information like the currently
@@ -517,7 +531,7 @@ export function makeServer<E = unknown>(options: ServerOptions<E>): Server<E> {
           case MessageType.Subscribe: {
             if (!ctx.acknowledged) return socket.close(4401, 'Unauthorized');
 
-            const { id } = message;
+            const { id, payload } = message;
             if (id in ctx.subscriptions)
               return socket.close(4409, `Subscriber for ${id} already exists`);
 
@@ -593,12 +607,17 @@ export function makeServer<E = unknown>(options: ServerOptions<E>): Server<E> {
               if (!schema)
                 throw new Error('The GraphQL schema is not provided');
 
-              const { operationName, query, variables } = message.payload;
+              const args = {
+                operationName: payload.operationName,
+                document: parse(payload.query),
+                variableValues: payload.variables,
+              };
               execArgs = {
-                schema,
-                operationName,
-                document: parse(query),
-                variableValues: variables,
+                ...args,
+                schema:
+                  typeof schema === 'function'
+                    ? await schema(ctx, message, args)
+                    : schema,
               };
               const validationErrors = validate(
                 execArgs.schema,

src/tests/fixtures/simple.ts

@@ -5,6 +5,7 @@ import {
   execute,
   subscribe,
   GraphQLNonNull,
+  GraphQLSchemaConfig,
 } from 'graphql';
 import { EventEmitter } from 'events';
 import WebSocket from 'ws';
@@ -57,7 +58,7 @@ function pong(key = 'global'): void {
   }
 }
 
-export const schema = new GraphQLSchema({
+export const schemaConfig: GraphQLSchemaConfig = {
   query: new GraphQLObjectType({
     name: 'Query',
     fields: {
@@ -116,7 +117,9 @@ export const schema = new GraphQLSchema({
       },
     },
   }),
-});
+};
+
+export const schema = new GraphQLSchema(schemaConfig);
 
 export async function startTServer(
   options: Partial<ServerOptions<Extra>> = {},

src/tests/server.ts

@@ -6,10 +6,11 @@ import {
   GraphQLError,
   ExecutionArgs,
   ExecutionResult,
+  GraphQLSchema,
 } from 'graphql';
 import { GRAPHQL_TRANSPORT_WS_PROTOCOL } from '../protocol';
 import { MessageType, parseMessage, stringifyMessage } from '../message';
-import { schema, startTServer } from './fixtures/simple';
+import { schema, schemaConfig, startTServer } from './fixtures/simple';
 import { createTClient } from './utils';
 
 /**
@@ -57,6 +58,41 @@ it('should allow connections with valid protocols only', async () => {
   console.warn = warn;
 });
 
+it('should use the schema resolved from a promise on subscribe', async (done) => {
+  expect.assertions(2);
+
+  const schema = new GraphQLSchema(schemaConfig);
+
+  const { url } = await startTServer({
+    schema: (_, msg) => {
+      expect(msg.id).toBe('1');
+      return Promise.resolve(schema);
+    },
+    execute: (args) => {
+      expect(args.schema).toBe(schema);
+      return execute(args);
+    },
+    onComplete: () => done(),
+  });
+  const client = await createTClient(url, GRAPHQL_TRANSPORT_WS_PROTOCOL);
+  client.ws.send(
+    stringifyMessage<MessageType.ConnectionInit>({
+      type: MessageType.ConnectionInit,
+    }),
+  );
+  await client.waitForMessage(); // ack
+
+  client.ws.send(
+    stringifyMessage<MessageType.Subscribe>({
+      id: '1',
+      type: MessageType.Subscribe,
+      payload: {
+        query: '{ getValue }',
+      },
+    }),
+  );
+});
+
 it('should use the provided roots as resolvers', async () => {
   const schema = buildSchema(`
     type Query {