Allows to add schema definition missing in the original schema

Author: IvanGoncharov

src/utilities/__tests__/extendSchema-test.js

@@ -1287,7 +1287,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
@@ -1302,6 +1302,31 @@ describe('extendSchema', () => {
       );
     });
 
+    it('adds schema definition missing in the original schema', () => {
+      let schema = new GraphQLSchema({
+        directives: [
+          new GraphQLDirective({
+            name: 'onSchema',
+            locations: ['SCHEMA'],
+          }),
+        ],
+      });
+      expect(schema.getQueryType()).to.equal(undefined);
+
+      const ast = parse(`
+        schema @onSchema {
+          query: Foo
+        }
+
+        type Foo {
+          bar: String
+        }
+      `);
+      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 { Kind } from '../language/kinds';
 import type {
   DocumentNode,
   DirectiveDefinitionNode,
+  SchemaDefinitionNode,
   SchemaExtensionNode,
 } from '../language/ast';
 
@@ -118,12 +119,16 @@ 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:
+        schemaDef = def;
+        break;
       case Kind.SCHEMA_EXTENSION:
         schemaExtensions.push(def);
         break;
@@ -190,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;
   }
@@ -222,9 +228,9 @@ export function extendSchema(
     subscription: extendMaybeNamedType(schema.getSubscriptionType()),
   };
 
-  // Then, incorporate all schema extensions.
-  for (const schemaExtension of schemaExtensions) {
-    if (schemaExtension.operationTypes) {
+  // Then, incorporate schema definition and all schema extensions.
+  for (const schemaExtension of [schemaDef, ...schemaExtensions]) {
+    if (schemaExtension && schemaExtension.operationTypes) {
       for (const operationType of schemaExtension.operationTypes) {
         const operation = operationType.operation;
         if (operationTypes[operation]) {

src/validation/rules/LoneSchemaDefinition.js

@@ -27,10 +27,18 @@ export function canNotDefineSchemaWithinExtension(): string {
 export function LoneSchemaDefinition(
   context: SDLValidationContext,
 ): ASTVisitor {
+  const oldSchema = context.getSchema();
+  const alreadyDefined =
+    oldSchema &&
+    (oldSchema.astNode ||
+      oldSchema.getQueryType() ||
+      oldSchema.getMutationType() ||
+      oldSchema.getSubscriptionType());
+
   const schemaNodes = [];
   return {
     SchemaDefinition(node) {
-      if (context.getSchema()) {
+      if (alreadyDefined) {
         context.reportError(
           new GraphQLError(canNotDefineSchemaWithinExtension(), [node]),
         );