Fix extend schema from #1314

src/utilities/__tests__/extendSchema-test.js

@@ -987,14 +987,25 @@ describe('extendSchema', () => {
           hearSomething: String
         }
       `);
-      const schema = extendSchema(testSchema, ast);
-      expect(schema.extensionASTNodes.map(print).join('\n')).to.equal(dedent`
+      let schema = extendSchema(testSchema, ast);
+
+      const secondAST = parse(`
+        extend schema @foo
+
+        directive @foo on SCHEMA
+      `);
+      schema = extendSchema(schema, secondAST);
+
+      const nodes = schema.extensionASTNodes;
+      expect(nodes.map(n => print(n) + '\n').join('')).to.equal(dedent`
         extend schema {
           mutation: Mutation
         }
         extend schema {
           subscription: Subscription
-        }`);
+        }
+        extend schema @foo
+      `);
     });
 
     it('does not allow redefining an existing root type', () => {

src/utilities/extendSchema.js

@@ -199,17 +199,10 @@ export function extendSchema(
   const extendTypeCache = Object.create(null);
 
   // Get the extended root operation types.
-  const existingQueryType = schema.getQueryType();
-  const existingMutationType = schema.getMutationType();
-  const existingSubscriptionType = schema.getSubscriptionType();
   const operationTypes = {
-    query: existingQueryType ? getExtendedType(existingQueryType) : null,
-    mutation: existingMutationType
-      ? getExtendedType(existingMutationType)
-      : null,
-    subscription: existingSubscriptionType
-      ? getExtendedType(existingSubscriptionType)
-      : null,
+    query: getExtendedMaybeType(schema.getQueryType()),
+    mutation: getExtendedMaybeType(schema.getMutationType()),
+    subscription: getExtendedMaybeType(schema.getSubscriptionType()),
   };
 
   // Then, incorporate all schema extensions.
@@ -220,11 +213,21 @@ export function extendSchema(
         if (operationTypes[operation]) {
           throw new Error(`Must provide only one ${operation} type in schema.`);
         }
-        operationTypes[operation] = astBuilder.buildType(operationType.type);
+        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);
       });
     }
   });
 
+  const schemaExtensionASTNodes = schemaExtensions
+    ? schema.extensionASTNodes
+      ? schema.extensionASTNodes.concat(schemaExtensions)
+      : schemaExtensions
+    : schema.extensionASTNodes;
+
   const types = [
     // Iterate through all types, getting the type definition for each, ensuring
     // that any type not directly referenced by a field will get created.
@@ -249,7 +252,7 @@ export function extendSchema(
     types,
     directives: getMergedDirectives(),
     astNode: schema.astNode,
-    extensionASTNodes: schemaExtensions,
+    extensionASTNodes: schemaExtensionASTNodes,
     allowedLegacyNames,
   });
 
@@ -265,6 +268,10 @@ export function extendSchema(
     );
   }
 
+  function getExtendedMaybeType<T: GraphQLNamedType>(type: ?T): ?T {
+    return type ? getExtendedType(type) : null;
+  }
+
   function getExtendedType<T: GraphQLNamedType>(type: T): T {
     if (!extendTypeCache[type.name]) {
       extendTypeCache[type.name] = extendType(type);