buildSchema/extendSchema: add support for extensions

IvanGoncharov · IvanGoncharov · commit cae97945312a · 2019-11-06T01:58:43.000+02:00
Fixes #922
diff --git a/src/utilities/__tests__/buildASTSchema-test.js b/src/utilities/__tests__/buildASTSchema-test.js
@@ -35,7 +35,7 @@ import {
 
 import { graphqlSync } from '../../graphql';
 
-import { printSchema } from '../schemaPrinter';
+import { printType, printSchema } from '../schemaPrinter';
 import { buildASTSchema, buildSchema } from '../buildASTSchema';
 
 /**
@@ -55,6 +55,17 @@ function printASTNode(obj) {
   return print(obj.astNode);
 }
 
+function printAllASTNodes(obj) {
+  invariant(obj != null);
+  invariant(obj.astNode != null);
+  invariant(obj.extensionASTNodes != null);
+
+  return print({
+    kind: Kind.DOCUMENT,
+    definitions: [obj.astNode, ...obj.extensionASTNodes],
+  });
+}
+
 describe('Schema Builder', () => {
   it('can use built schema for limited execution', () => {
     const schema = buildASTSchema(
@@ -737,6 +748,89 @@ describe('Schema Builder', () => {
     });
   });
 
+  it('Correctly extend scalar type', () => {
+    const scalarSDL = dedent`
+      scalar SomeScalar
+
+      extend scalar SomeScalar @foo
+
+      extend scalar SomeScalar @bar
+    `;
+    const schema = buildSchema(`
+      ${scalarSDL}
+      directive @foo on SCALAR
+      directive @bar on SCALAR
+    `);
+
+    const someScalar = assertScalarType(schema.getType('SomeScalar'));
+    expect(printType(someScalar) + '\n').to.equal(dedent`
+      scalar SomeScalar
+    `);
+
+    expect(printAllASTNodes(someScalar)).to.equal(scalarSDL);
+  });
+
+  it('Correctly extend object type', () => {
+    const objectSDL = dedent`
+      type SomeObject implements Foo {
+        first: String
+      }
+
+      extend type SomeObject implements Bar {
+        second: Int
+      }
+
+      extend type SomeObject implements Baz {
+        third: Float
+      }
+    `;
+    const schema = buildSchema(`
+      ${objectSDL}
+      interface Foo
+      interface Bar
+      interface Baz
+    `);
+
+    const someObject = assertObjectType(schema.getType('SomeObject'));
+    expect(printType(someObject) + '\n').to.equal(dedent`
+      type SomeObject implements Foo & Bar & Baz {
+        first: String
+        second: Int
+        third: Float
+      }
+    `);
+
+    expect(printAllASTNodes(someObject)).to.equal(objectSDL);
+  });
+
+  it('Correctly extend interface type', () => {
+    const interfaceSDL = dedent`
+      interface SomeInterface {
+        first: String
+      }
+
+      extend interface SomeInterface {
+        second: Int
+      }
+
+      extend interface SomeInterface {
+        third: Float
+      }
+    `;
+    const schema = buildSchema(interfaceSDL);
+
+    const someInterface = assertInterfaceType(schema.getType('SomeInterface'));
+    expect(printType(someInterface) + '\n').to.equal(dedent`
+      interface SomeInterface {
+        first: String
+        second: Int
+        third: Float
+      }
+    `);
+
+    expect(printAllASTNodes(someInterface)).to.equal(interfaceSDL);
+  });
+
   it('Correctly assign AST nodes', () => {
     const sdl = dedent`
       schema {
diff --git a/src/utilities/buildASTSchema.js b/src/utilities/buildASTSchema.js
@@ -6,20 +6,24 @@ import keyMap from '../jsutils/keyMap';
 import inspect from '../jsutils/inspect';
 import invariant from '../jsutils/invariant';
 import devAssert from '../jsutils/devAssert';
-import { type ObjMap } from '../jsutils/ObjMap';
+import { type ObjMap, type ReadOnlyObjMap } from '../jsutils/ObjMap';
 
 import { Kind } from '../language/kinds';
 import { type Source } from '../language/source';
 import { TokenKind } from '../language/tokenKind';
 import { type ParseOptions, parse } from '../language/parser';
-import { isTypeDefinitionNode } from '../language/predicates';
 import { dedentBlockStringValue } from '../language/blockString';
 import { type DirectiveLocationEnum } from '../language/directiveLocation';
+import {
+  isTypeDefinitionNode,
+  isTypeExtensionNode,
+} from '../language/predicates';
 import {
   type Location,
   type StringValueNode,
   type DocumentNode,
   type TypeNode,
+  type TypeExtensionNode,
   type NamedTypeNode,
   type SchemaDefinitionNode,
   type SchemaExtensionNode,
@@ -127,15 +131,26 @@ export function buildASTSchema(
     assertValidSDL(documentAST);
   }
 
+  // Collect the definitions and extensions found in the document.
   let schemaDef: ?SchemaDefinitionNode;
+  const schemaExts: Array<SchemaExtensionNode> = [];
   const typeDefs: Array<TypeDefinitionNode> = [];
+  const typeExtsMap: ObjMap<Array<TypeExtensionNode>> = Object.create(null);
   const directiveDefs: Array<DirectiveDefinitionNode> = [];
 
   for (const def of documentAST.definitions) {
     if (def.kind === Kind.SCHEMA_DEFINITION) {
       schemaDef = def;
+    } else if (def.kind === Kind.SCHEMA_EXTENSION) {
+      schemaExts.push(def);
     } else if (isTypeDefinitionNode(def)) {
       typeDefs.push(def);
+    } else if (isTypeExtensionNode(def)) {
+      const extendedTypeName = def.name.value;
+      const existingTypeExts = typeExtsMap[extendedTypeName];
+      typeExtsMap[extendedTypeName] = existingTypeExts
+        ? existingTypeExts.concat([def])
+        : [def];
     } else if (def.kind === Kind.DIRECTIVE_DEFINITION) {
       directiveDefs.push(def);
     }
@@ -149,7 +164,7 @@ export function buildASTSchema(
     return type;
   });
 
-  const typeMap = astBuilder.buildTypeMap(typeDefs);
+  const typeMap = astBuilder.buildTypeMap(typeDefs, typeExtsMap);
   const operationTypes = schemaDef
     ? astBuilder.getOperationTypes([schemaDef])
     : {
@@ -394,63 +409,97 @@ export class ASTDefinitionBuilder {
 
   buildTypeMap(
     nodes: $ReadOnlyArray<TypeDefinitionNode>,
+    extensionMap: ReadOnlyObjMap<$ReadOnlyArray<TypeExtensionNode>>,
   ): ObjMap<GraphQLNamedType> {
     const typeMap = Object.create(null);
     for (const node of nodes) {
       const name = node.name.value;
-      typeMap[name] = stdTypeMap[name] || this._buildType(node);
+      typeMap[name] =
+        stdTypeMap[name] || this._buildType(node, extensionMap[name] || []);
     }
     return typeMap;
   }
 
-  _buildType(astNode: TypeDefinitionNode): GraphQLNamedType {
+  _buildType(
+    astNode: TypeDefinitionNode,
+    extensionNodes: $ReadOnlyArray<TypeExtensionNode>,
+  ): GraphQLNamedType {
     const name = astNode.name.value;
     const description = getDescription(astNode, this._options);
 
     switch (astNode.kind) {
-      case Kind.OBJECT_TYPE_DEFINITION:
+      case Kind.OBJECT_TYPE_DEFINITION: {
+        const extensionASTNodes = (extensionNodes: any);
+        const allNodes = [astNode, ...extensionASTNodes];
+
         return new GraphQLObjectType({
           name,
           description,
-          interfaces: () => this.buildInterfaces([astNode]),
-          fields: () => this.buildFieldMap([astNode]),
+          interfaces: () => this.buildInterfaces(allNodes),
+          fields: () => this.buildFieldMap(allNodes),
           astNode,
+          extensionASTNodes,
         });
-      case Kind.INTERFACE_TYPE_DEFINITION:
+      }
+      case Kind.INTERFACE_TYPE_DEFINITION: {
+        const extensionASTNodes = (extensionNodes: any);
+        const allNodes = [astNode, ...extensionASTNodes];
+
         return new GraphQLInterfaceType({
           name,
           description,
-          interfaces: () => this.buildInterfaces([astNode]),
-          fields: () => this.buildFieldMap([astNode]),
+          interfaces: () => this.buildInterfaces(allNodes),
+          fields: () => this.buildFieldMap(allNodes),
           astNode,
+          extensionASTNodes,
         });
-      case Kind.ENUM_TYPE_DEFINITION:
+      }
+      case Kind.ENUM_TYPE_DEFINITION: {
+        const extensionASTNodes = (extensionNodes: any);
+        const allNodes = [astNode, ...extensionASTNodes];
+
         return new GraphQLEnumType({
           name,
           description,
-          values: this.buildEnumValueMap([astNode]),
+          values: this.buildEnumValueMap(allNodes),
           astNode,
+          extensionASTNodes,
         });
-      case Kind.UNION_TYPE_DEFINITION:
+      }
+      case Kind.UNION_TYPE_DEFINITION: {
+        const extensionASTNodes = (extensionNodes: any);
+        const allNodes = [astNode, ...extensionASTNodes];
+
         return new GraphQLUnionType({
           name,
           description,
-          types: () => this.buildUnionTypes([astNode]),
+          types: () => this.buildUnionTypes(allNodes),
           astNode,
+          extensionASTNodes,
         });
-      case Kind.SCALAR_TYPE_DEFINITION:
+      }
+      case Kind.SCALAR_TYPE_DEFINITION: {
+        const extensionASTNodes = (extensionNodes: any);
+
         return new GraphQLScalarType({
           name,
           description,
           astNode,
+          extensionASTNodes,
         });
-      case Kind.INPUT_OBJECT_TYPE_DEFINITION:
+      }
+      case Kind.INPUT_OBJECT_TYPE_DEFINITION: {
+        const extensionASTNodes = (extensionNodes: any);
+        const allNodes = [astNode, ...extensionASTNodes];
+
         return new GraphQLInputObjectType({
           name,
           description,
-          fields: () => this.buildInputFieldMap([astNode]),
+          fields: () => this.buildInputFieldMap(allNodes),
           astNode,
+          extensionASTNodes,
         });
+      }
     }
 
     // Not reachable. All possible type definition nodes have been considered.
diff --git a/src/utilities/extendSchema.js b/src/utilities/extendSchema.js
@@ -156,7 +156,7 @@ export function extendSchema(
     return type;
   });
 
-  const typeMap = astBuilder.buildTypeMap(typeDefs);
+  const typeMap = astBuilder.buildTypeMap(typeDefs, typeExtsMap);
   const schemaConfig = schema.toConfig();
   for (const existingType of schemaConfig.types) {
     typeMap[existingType.name] = extendNamedType(existingType);
