Support extendSchema()

Author: leebyron

src/language/ast.js

@@ -396,7 +396,7 @@ export type TypeSystemDefinitionNode =
 export type SchemaDefinitionNode = {
   +kind: 'SchemaDefinition',
   +loc?: Location,
-  +directives: $ReadOnlyArray<DirectiveNode>,
+  +directives?: $ReadOnlyArray<DirectiveNode>,
   +operationTypes: $ReadOnlyArray<OperationTypeDefinitionNode>,
 };
 
@@ -517,8 +517,8 @@ export type TypeSystemExtensionNode = SchemaExtensionNode | TypeExtensionNode;
 export type SchemaExtensionNode = {
   +kind: 'SchemaExtension',
   +loc?: Location,
-  +directives: $ReadOnlyArray<DirectiveNode>,
-  +operationTypes: $ReadOnlyArray<OperationTypeDefinitionNode>,
+  +directives?: $ReadOnlyArray<DirectiveNode>,
+  +operationTypes?: $ReadOnlyArray<OperationTypeDefinitionNode>,
 };
 
 // Type Extensions

src/type/schema.js

@@ -21,7 +21,10 @@ import type {
   GraphQLAbstractType,
   GraphQLObjectType,
 } from './definition';
-import type { SchemaDefinitionNode } from '../language/ast';
+import type {
+  SchemaDefinitionNode,
+  SchemaExtensionNode,
+} from '../language/ast';
 import {
   GraphQLDirective,
   isDirective,
@@ -73,6 +76,7 @@ export function isSchema(schema) {
  */
 export class GraphQLSchema {
   astNode: ?SchemaDefinitionNode;
+  extensionASTNodes: ?$ReadOnlyArray<SchemaExtensionNode>;
   _queryType: ?GraphQLObjectType;
   _mutationType: ?GraphQLObjectType;
   _subscriptionType: ?GraphQLObjectType;
@@ -120,6 +124,7 @@ export class GraphQLSchema {
     // Provide specified directives (e.g. @include and @skip) by default.
     this._directives = config.directives || specifiedDirectives;
     this.astNode = config.astNode;
+    this.extensionASTNodes = config.extensionASTNodes;
 
     // Build type map now to detect any errors within this schema.
     let initialTypes: Array<?GraphQLNamedType> = [
@@ -255,6 +260,7 @@ export type GraphQLSchemaConfig = {
   types?: ?Array<GraphQLNamedType>,
   directives?: ?Array<GraphQLDirective>,
   astNode?: ?SchemaDefinitionNode,
+  extensionASTNodes?: ?$ReadOnlyArray<SchemaExtensionNode>,
   ...GraphQLSchemaValidationOptions,
 };
 

src/utilities/__tests__/extendSchema-test.js

@@ -876,4 +876,182 @@ describe('extendSchema', () => {
       );
     });
   });
+
+  describe('can add additional root operation types', () => {
+    it('does not automatically include common root type names', () => {
+      const ast = parse(`
+        type Mutation {
+          doSomething: String
+        }
+      `);
+      const schema = extendSchema(testSchema, ast);
+      expect(schema.getMutationType()).to.equal(null);
+    });
+
+    it('does not allow new schema within an extension', () => {
+      const ast = parse(`
+        schema {
+          mutation: Mutation
+        }
+
+        type Mutation {
+          doSomething: String
+        }
+      `);
+      expect(() => extendSchema(testSchema, ast)).to.throw(
+        'Cannot define a new schema within a schema extension.',
+      );
+    });
+
+    it('adds new root types via schema extension', () => {
+      const ast = parse(`
+        extend schema {
+          mutation: Mutation
+        }
+
+        type Mutation {
+          doSomething: String
+        }
+      `);
+      const schema = extendSchema(testSchema, ast);
+      const mutationType = schema.getMutationType();
+      expect(mutationType && mutationType.name).to.equal('Mutation');
+    });
+
+    it('adds multiple new root types via schema extension', () => {
+      const ast = parse(`
+        extend schema {
+          mutation: Mutation
+          subscription: Subscription
+        }
+
+        type Mutation {
+          doSomething: String
+        }
+
+        type Subscription {
+          hearSomething: String
+        }
+      `);
+      const schema = extendSchema(testSchema, ast);
+      const mutationType = schema.getMutationType();
+      const subscriptionType = schema.getSubscriptionType();
+      expect(mutationType && mutationType.name).to.equal('Mutation');
+      expect(subscriptionType && subscriptionType.name).to.equal(
+        'Subscription',
+      );
+    });
+
+    it('applies multiple schema extensions', () => {
+      const ast = parse(`
+        extend schema {
+          mutation: Mutation
+        }
+
+        extend schema {
+          subscription: Subscription
+        }
+
+        type Mutation {
+          doSomething: String
+        }
+
+        type Subscription {
+          hearSomething: String
+        }
+      `);
+      const schema = extendSchema(testSchema, ast);
+      const mutationType = schema.getMutationType();
+      const subscriptionType = schema.getSubscriptionType();
+      expect(mutationType && mutationType.name).to.equal('Mutation');
+      expect(subscriptionType && subscriptionType.name).to.equal(
+        'Subscription',
+      );
+    });
+
+    it('schema extension AST are available from schema object', () => {
+      const ast = parse(`
+        extend schema {
+          mutation: Mutation
+        }
+
+        extend schema {
+          subscription: Subscription
+        }
+
+        type Mutation {
+          doSomething: String
+        }
+
+        type Subscription {
+          hearSomething: String
+        }
+      `);
+      const schema = extendSchema(testSchema, ast);
+      expect(schema.extensionASTNodes.map(print).join('\n')).to.equal(dedent`
+        extend schema {
+          mutation: Mutation
+        }
+        extend schema {
+          subscription: Subscription
+        }`);
+    });
+
+    it('does not allow redefining an existing root type', () => {
+      const ast = parse(`
+        extend schema {
+          query: SomeType
+        }
+
+        type SomeType {
+          seeSomething: String
+        }
+      `);
+      expect(() => extendSchema(testSchema, ast)).to.throw(
+        'Must provide only one query type in schema.',
+      );
+    });
+
+    it('does not allow defining a root operation type twice', () => {
+      const ast = parse(`
+        extend schema {
+          mutation: Mutation
+        }
+
+        extend schema {
+          mutation: Mutation
+        }
+
+        type Mutation {
+          doSomething: String
+        }
+      `);
+      expect(() => extendSchema(testSchema, ast)).to.throw(
+        'Must provide only one mutation type in schema.',
+      );
+    });
+
+    it('does not allow defining a root operation type with different types', () => {
+      const ast = parse(`
+        extend schema {
+          mutation: Mutation
+        }
+
+        extend schema {
+          mutation: SomethingElse
+        }
+
+        type Mutation {
+          doSomething: String
+        }
+
+        type SomethingElse {
+          doSomethingElse: String
+        }
+      `);
+      expect(() => extendSchema(testSchema, ast)).to.throw(
+        'Must provide only one mutation type in schema.',
+      );
+    });
+  });
 });

src/utilities/extendSchema.js

@@ -35,7 +35,11 @@ import { GraphQLDirective } from '../type/directives';
 import { Kind } from '../language/kinds';
 
 import type { GraphQLType, GraphQLNamedType } from '../type/definition';
-import type { DocumentNode, DirectiveDefinitionNode } from '../language/ast';
+import type {
+  DocumentNode,
+  DirectiveDefinitionNode,
+  SchemaExtensionNode,
+} from '../language/ast';
 
 type Options = {|
   ...GraphQLSchemaValidationOptions,
@@ -88,9 +92,21 @@ export function extendSchema(
   // have the same name. For example, a type named "skip".
   const directiveDefinitions: Array<DirectiveDefinitionNode> = [];
 
+  // Schema extensions are collected which may add additional operation types.
+  const schemaExtensions: Array<SchemaExtensionNode> = [];
+
   for (let i = 0; i < documentAST.definitions.length; i++) {
     const def = documentAST.definitions[i];
     switch (def.kind) {
+      case Kind.SCHEMA_DEFINITION:
+        // Sanity check that a schema extension is not defining a new schema
+        throw new GraphQLError(
+          'Cannot define a new schema within a schema extension.',
+          [def],
+        );
+      case Kind.SCHEMA_EXTENSION:
+        schemaExtensions.push(def);
+        break;
       case Kind.OBJECT_TYPE_DEFINITION:
       case Kind.INTERFACE_TYPE_DEFINITION:
       case Kind.ENUM_TYPE_DEFINITION:
@@ -156,7 +172,8 @@ export function extendSchema(
   if (
     Object.keys(typeExtensionsMap).length === 0 &&
     Object.keys(typeDefinitionMap).length === 0 &&
-    directiveDefinitions.length === 0
+    directiveDefinitions.length === 0 &&
+    schemaExtensions.length === 0
   ) {
     return schema;
   }
@@ -181,21 +198,32 @@ export function extendSchema(
 
   const extendTypeCache = Object.create(null);
 
-  // Get the root Query, Mutation, and Subscription object types.
+  // Get the extended root operation types.
   const existingQueryType = schema.getQueryType();
-  const queryType = existingQueryType
-    ? getExtendedType(existingQueryType)
-    : null;
-
   const existingMutationType = schema.getMutationType();
-  const mutationType = existingMutationType
-    ? getExtendedType(existingMutationType)
-    : null;
-
   const existingSubscriptionType = schema.getSubscriptionType();
-  const subscriptionType = existingSubscriptionType
-    ? getExtendedType(existingSubscriptionType)
-    : null;
+  const operationTypes = {
+    query: existingQueryType ? getExtendedType(existingQueryType) : null,
+    mutation: existingMutationType
+      ? getExtendedType(existingMutationType)
+      : null,
+    subscription: existingSubscriptionType
+      ? getExtendedType(existingSubscriptionType)
+      : null,
+  };
+
+  // Then, incorporate all schema extensions.
+  schemaExtensions.forEach(schemaExtension => {
+    if (schemaExtension.operationTypes) {
+      schemaExtension.operationTypes.forEach(operationType => {
+        const operation = operationType.operation;
+        if (operationTypes[operation]) {
+          throw new Error(`Must provide only one ${operation} type in schema.`);
+        }
+        operationTypes[operation] = astBuilder.buildType(operationType.type);
+      });
+    }
+  });
 
   const types = [
     // Iterate through all types, getting the type definition for each, ensuring
@@ -215,12 +243,13 @@ export function extendSchema(
 
   // Then produce and return a Schema with these types.
   return new GraphQLSchema({
-    query: queryType,
-    mutation: mutationType,
-    subscription: subscriptionType,
+    query: operationTypes.query,
+    mutation: operationTypes.mutation,
+    subscription: operationTypes.subscription,
     types,
     directives: getMergedDirectives(),
     astNode: schema.astNode,
+    extensionASTNodes: schemaExtensions,
     allowedLegacyNames,
   });