Fix extend schema from #1314 (#1323)

* Simplify + fix Flow types

* Fix 'extensionASTNodes' for schema extended twice

* Replace 'getAll*Nodes' with generic 'getAllNodes' function

* Correctly report error for invalid root type in extended schema

Author: IvanGoncharov

src/type/__tests__/validation-test.js

@@ -323,6 +323,59 @@ describe('Type System: A Schema must have Object root types', () => {
     ]);
   });
 
+  it('rejects a schema extended with invalid root types', () => {
+    let schema = buildSchema(`
+      input SomeInputObject {
+        test: String
+      }
+    `);
+
+    schema = extendSchema(
+      schema,
+      parse(`
+        extend schema {
+          query: SomeInputObject
+        }
+      `),
+    );
+
+    schema = extendSchema(
+      schema,
+      parse(`
+        extend schema {
+          mutation: SomeInputObject
+        }
+      `),
+    );
+
+    schema = extendSchema(
+      schema,
+      parse(`
+        extend schema {
+          subscription: SomeInputObject
+        }
+      `),
+    );
+
+    expect(validateSchema(schema)).to.deep.equal([
+      {
+        message:
+          'Query root type must be Object type, it cannot be SomeInputObject.',
+        locations: [{ line: 3, column: 18 }],
+      },
+      {
+        message:
+          'Mutation root type must be Object type if provided, it cannot be SomeInputObject.',
+        locations: [{ line: 3, column: 21 }],
+      },
+      {
+        message:
+          'Subscription root type must be Object type if provided, it cannot be SomeInputObject.',
+        locations: [{ line: 3, column: 25 }],
+      },
+    ]);
+  });
+
   it('rejects a Schema whose directives are incorrectly typed', () => {
     const schema = new GraphQLSchema({
       query: SomeObjectType,

src/type/validate.js

@@ -36,10 +36,6 @@ import objectValues from '../jsutils/objectValues';
 import { GraphQLError } from '../error/GraphQLError';
 import type {
   ASTNode,
-  ObjectTypeDefinitionNode,
-  ObjectTypeExtensionNode,
-  InterfaceTypeDefinitionNode,
-  InterfaceTypeExtensionNode,
   FieldDefinitionNode,
   EnumValueDefinitionNode,
   InputValueDefinitionNode,
@@ -156,13 +152,17 @@ function getOperationTypeNode(
   type: GraphQLObjectType,
   operation: string,
 ): ?ASTNode {
-  const astNode = schema.astNode;
-  const operationTypeNode =
-    astNode &&
-    astNode.operationTypes.find(
-      operationType => operationType.operation === operation,
-    );
-  return operationTypeNode ? operationTypeNode.type : type && type.astNode;
+  for (const node of getAllNodes(schema)) {
+    if (node.operationTypes) {
+      for (const operationType of node.operationTypes) {
+        if (operationType.operation === operation) {
+          return operationType.type;
+        }
+      }
+    }
+  }
+
+  return type.astNode;
 }
 
 function validateDirectives(context: SchemaValidationContext): void {
@@ -283,7 +283,7 @@ function validateFields(
   if (fields.length === 0) {
     context.reportError(
       `Type ${type.name} must define one or more fields.`,
-      getAllObjectOrInterfaceNodes(type),
+      getAllNodes(type),
     );
   }
 
@@ -578,29 +578,16 @@ function validateInputFields(
   });
 }
 
-function getAllObjectNodes(
-  type: GraphQLObjectType,
-): $ReadOnlyArray<ObjectTypeDefinitionNode | ObjectTypeExtensionNode> {
-  return type.astNode
-    ? type.extensionASTNodes
-      ? [type.astNode].concat(type.extensionASTNodes)
-      : [type.astNode]
-    : type.extensionASTNodes || [];
-}
-
-function getAllObjectOrInterfaceNodes(
-  type: GraphQLObjectType | GraphQLInterfaceType,
-): $ReadOnlyArray<
-  | ObjectTypeDefinitionNode
-  | ObjectTypeExtensionNode
-  | InterfaceTypeDefinitionNode
-  | InterfaceTypeExtensionNode,
-> {
-  return type.astNode
-    ? type.extensionASTNodes
-      ? [type.astNode].concat(type.extensionASTNodes)
-      : [type.astNode]
-    : type.extensionASTNodes || [];
+function getAllNodes<T: ASTNode, K: ASTNode>(object: {
+  +astNode: ?T,
+  +extensionASTNodes?: ?$ReadOnlyArray<K>,
+}): $ReadOnlyArray<T | K> {
+  const { astNode, extensionASTNodes } = object;
+  return astNode
+    ? extensionASTNodes
+      ? [astNode].concat(extensionASTNodes)
+      : [astNode]
+    : extensionASTNodes || [];
 }
 
 function getImplementsInterfaceNode(
@@ -615,7 +602,7 @@ function getAllImplementsInterfaceNodes(
   iface: GraphQLInterfaceType,
 ): $ReadOnlyArray<NamedTypeNode> {
   const implementsNodes = [];
-  const astNodes = getAllObjectNodes(type);
+  const astNodes = getAllNodes(type);
   for (let i = 0; i < astNodes.length; i++) {
     const astNode = astNodes[i];
     if (astNode && astNode.interfaces) {
@@ -641,7 +628,7 @@ function getAllFieldNodes(
   fieldName: string,
 ): $ReadOnlyArray<FieldDefinitionNode> {
   const fieldNodes = [];
-  const astNodes = getAllObjectOrInterfaceNodes(type);
+  const astNodes = getAllNodes(type);
   for (let i = 0; i < astNodes.length; i++) {
     const astNode = astNodes[i];
     if (astNode && astNode.fields) {

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);