Allows to add schema definition missing in the original schema

src/utilities/__tests__/extendSchema-test.js

@@ -93,6 +93,26 @@ const SomeInputType = new GraphQLInputObjectType({
   }),
 });
 
+const FooDirective = new GraphQLDirective({
+  name: 'foo',
+  args: {
+    input: { type: SomeInputType },
+  },
+  locations: [
+    DirectiveLocation.SCHEMA,
+    DirectiveLocation.SCALAR,
+    DirectiveLocation.OBJECT,
+    DirectiveLocation.FIELD_DEFINITION,
+    DirectiveLocation.ARGUMENT_DEFINITION,
+    DirectiveLocation.INTERFACE,
+    DirectiveLocation.UNION,
+    DirectiveLocation.ENUM,
+    DirectiveLocation.ENUM_VALUE,
+    DirectiveLocation.INPUT_OBJECT,
+    DirectiveLocation.INPUT_FIELD_DEFINITION,
+  ],
+});
+
 const testSchema = new GraphQLSchema({
   query: new GraphQLObjectType({
     name: 'Query',
@@ -112,27 +132,7 @@ const testSchema = new GraphQLSchema({
     }),
   }),
   types: [FooType, BarType],
-  directives: specifiedDirectives.concat([
-    new GraphQLDirective({
-      name: 'foo',
-      args: {
-        input: { type: SomeInputType },
-      },
-      locations: [
-        DirectiveLocation.SCHEMA,
-        DirectiveLocation.SCALAR,
-        DirectiveLocation.OBJECT,
-        DirectiveLocation.FIELD_DEFINITION,
-        DirectiveLocation.ARGUMENT_DEFINITION,
-        DirectiveLocation.INTERFACE,
-        DirectiveLocation.UNION,
-        DirectiveLocation.ENUM,
-        DirectiveLocation.ENUM_VALUE,
-        DirectiveLocation.INPUT_OBJECT,
-        DirectiveLocation.INPUT_FIELD_DEFINITION,
-      ],
-    }),
-  ]),
+  directives: specifiedDirectives.concat([FooDirective]),
 });
 
 function extendTestSchema(sdl, options) {
@@ -1271,7 +1271,7 @@ describe('extendSchema', () => {
       expect(schema.getMutationType()).to.equal(null);
     });
 
-    it('does not allow new schema within an extension', () => {
+    it('does not allow overriding schema within an extension', () => {
       const sdl = `
         schema {
           mutation: Mutation
@@ -1286,6 +1286,23 @@ describe('extendSchema', () => {
       );
     });
 
+    it('adds schema definition missing in the original schema', () => {
+      let schema = new GraphQLSchema({
+        directives: [FooDirective],
+        types: [FooType],
+      });
+      expect(schema.getQueryType()).to.equal(undefined);
+
+      const ast = parse(`
+        schema @foo {
+          query: Foo
+        }
+      `);
+      schema = extendSchema(schema, ast);
+      const queryType = schema.getQueryType();
+      expect(queryType).to.have.property('name', 'Foo');
+    });
+
     it('adds new root types via schema extension', () => {
       const schema = extendTestSchema(`
         extend schema {

src/utilities/extendSchema.js

@@ -53,6 +53,7 @@ import type {
   DocumentNode,
   DirectiveDefinitionNode,
   SchemaExtensionNode,
+  SchemaDefinitionNode,
 } from '../language/ast';
 
 type Options = {|
@@ -106,18 +107,28 @@ export function extendSchema(
   // have the same name. For example, a type named "skip".
   const directiveDefinitions: Array<DirectiveDefinitionNode> = [];
 
+  let schemaDef: ?SchemaDefinitionNode;
   // 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],
-        );
+        // Sanity check that a schema extension is not overriding the schema
+        if (
+          schema.astNode ||
+          schema.getQueryType() ||
+          schema.getMutationType() ||
+          schema.getSubscriptionType()
+        ) {
+          throw new GraphQLError(
+            'Cannot define a new schema within a schema extension.',
+            [def],
+          );
+        }
+        schemaDef = def;
+        break;
       case Kind.SCHEMA_EXTENSION:
         schemaExtensions.push(def);
         break;
@@ -184,7 +195,8 @@ export function extendSchema(
     Object.keys(typeExtensionsMap).length === 0 &&
     Object.keys(typeDefinitionMap).length === 0 &&
     directiveDefinitions.length === 0 &&
-    schemaExtensions.length === 0
+    schemaExtensions.length === 0 &&
+    !schemaDef
   ) {
     return schema;
   }
@@ -216,19 +228,28 @@ export function extendSchema(
     subscription: extendMaybeNamedType(schema.getSubscriptionType()),
   };
 
-  // Then, incorporate all schema extensions.
+  if (schemaDef) {
+    for (const { operation, type } of schemaDef.operationTypes) {
+      if (operationTypes[operation]) {
+        throw new Error(`Must provide only one ${operation} type in schema.`);
+      }
+      // Note: While this could make early assertions to get the correctly
+      // typed values, that would throw immediately while type system
+      // validation with validateSchema() will produce more actionable results.
+      operationTypes[operation] = (astBuilder.buildType(type): any);
+    }
+  }
+  // Then, incorporate schema definition and all schema extensions.
   for (const schemaExtension of schemaExtensions) {
     if (schemaExtension.operationTypes) {
-      for (const operationType of schemaExtension.operationTypes) {
-        const operation = operationType.operation;
+      for (const { operation, type } of schemaExtension.operationTypes) {
         if (operationTypes[operation]) {
           throw new Error(`Must provide only one ${operation} type in schema.`);
         }
-        const typeRef = operationType.type;
         // Note: While this could make early assertions to get the correctly
         // typed values, that would throw immediately while type system
         // validation with validateSchema() will produce more actionable results.
-        operationTypes[operation] = (astBuilder.buildType(typeRef): any);
+        operationTypes[operation] = (astBuilder.buildType(type): any);
       }
     }
   }
@@ -254,9 +275,7 @@ export function extendSchema(
 
   // Then produce and return a Schema with these types.
   return new GraphQLSchema({
-    query: operationTypes.query,
-    mutation: operationTypes.mutation,
-    subscription: operationTypes.subscription,
+    ...operationTypes,
     types,
     directives: getMergedDirectives(),
     astNode: schema.astNode,