Make 'extendSchema' assign 'astNode' & 'extensionASTNodes'

src/utilities/__tests__/extendSchema-test.js

@@ -10,9 +10,10 @@
 import { describe, it } from 'mocha';
 import { expect } from 'chai';
 import dedent from '../../jsutils/dedent';
+import { buildSchema } from '../buildASTSchema';
 import { extendSchema } from '../extendSchema';
 import { execute } from '../../execution';
-import { parse } from '../../language';
+import { parse, print } from '../../language';
 import { printSchema } from '../schemaPrinter';
 import {
   GraphQLSchema,
@@ -196,6 +197,96 @@ describe('extendSchema', () => {
     `);
   });
 
+  it('correctly assign AST nodes to new and extended types', () => {
+    const schemaFromIDL = buildSchema(`
+      type Query {
+        dummyField: String
+      }
+    `);
+    const extensionAst = parse(`
+      extend type Query {
+        newField(testArg: TestInput): TestEnum
+      }
+
+      enum TestEnum {
+        TEST_VALUE
+      }
+
+      input TestInput {
+        testInputField: TestEnum
+      }
+    `);
+    const extendedSchema = extendSchema(schemaFromIDL, extensionAst);
+    const secondExtensionAst = parse(`
+      extend type Query {
+        oneMoreNewField: TestUnion
+      }
+
+      union TestUnion = TestType
+
+      interface TestInterface {
+        interfaceField: String
+      }
+
+      type TestType implements TestInterface {
+        interfaceField: String
+      }
+
+      directive @test(arg: Int) on FIELD
+    `);
+    const extendedTwiceSchema = extendSchema(extendedSchema,
+      secondExtensionAst);
+
+    const query = extendedTwiceSchema.getType('Query');
+    const testInput = extendedTwiceSchema.getType('TestInput');
+    const testEnum = extendedTwiceSchema.getType('TestEnum');
+    const testUnion = extendedTwiceSchema.getType('TestUnion');
+    const testInterface = extendedTwiceSchema.getType('TestInterface');
+    const testType = extendedTwiceSchema.getType('TestType');
+    const testDirective = extendedTwiceSchema.getDirective('test');
+
+    expect(query.extensionASTNodes).to.have.lengthOf(2);
+    expect(testType.extensionASTNodes).to.have.lengthOf(0);
+
+    const restoredExtensionAST = parse(
+      print(query.extensionASTNodes[0]) + '\n' +
+      print(query.extensionASTNodes[1]) + '\n' +
+      print(testInput.astNode) + '\n' +
+      print(testEnum.astNode) + '\n' +
+      print(testUnion.astNode) + '\n' +
+      print(testInterface.astNode) + '\n' +
+      print(testType.astNode) + '\n' +
+      print(testDirective.astNode)
+    );
+    expect(
+      printSchema(extendSchema(schemaFromIDL, restoredExtensionAST))
+    ).to.be.equal(printSchema(extendedTwiceSchema));
+
+    const newField = query.getFields().newField;
+    expect(print(newField.astNode)).to.equal(
+      'newField(testArg: TestInput): TestEnum'
+    );
+    expect(print(newField.args[0].astNode)).to.equal(
+      'testArg: TestInput'
+    );
+    expect(print(query.getFields().oneMoreNewField.astNode)).to.equal(
+      'oneMoreNewField: TestUnion'
+    );
+    expect(print(testInput.getFields().testInputField.astNode)).to.equal(
+      'testInputField: TestEnum'
+    );
+    expect(print(testEnum.getValue('TEST_VALUE').astNode)).to.equal(
+      'TEST_VALUE'
+    );
+    expect(print(testInterface.getFields().interfaceField.astNode)).to.equal(
+      'interfaceField: String'
+    );
+    expect(print(testType.getFields().interfaceField.astNode)).to.equal(
+      'interfaceField: String'
+    );
+    expect(print(testDirective.args[0].astNode)).to.equal('arg: Int');
+  });
+
   it('builds types with deprecated fields/values', () => {
     const ast = parse(`
       type TypeWithDeprecatedField {

src/utilities/extendSchema.js

@@ -240,6 +240,7 @@ export function extendSchema(
     subscription: subscriptionType,
     types,
     directives: getMergedDirectives(),
+    astNode: schema.astNode,
   });
 
   // Below are functions used for producing this schema that have closed over
@@ -332,11 +333,19 @@ export function extendSchema(
   }
 
   function extendObjectType(type: GraphQLObjectType): GraphQLObjectType {
+    const name = type.name;
+    let extensionASTNodes = type.extensionASTNodes;
+    if (typeExtensionsMap[name]) {
+      extensionASTNodes = extensionASTNodes.concat(typeExtensionsMap[name]);
+    }
+
     return new GraphQLObjectType({
-      name: type.name,
+      name,
       description: type.description,
       interfaces: () => extendImplementedInterfaces(type),
       fields: () => extendFieldMap(type),
+      astNode: type.astNode,
+      extensionASTNodes,
       isTypeOf: type.isTypeOf,
     });
   }
@@ -348,6 +357,7 @@ export function extendSchema(
       name: type.name,
       description: type.description,
       fields: () => extendFieldMap(type),
+      astNode: type.astNode,
       resolveType: type.resolveType,
     });
   }
@@ -357,6 +367,7 @@ export function extendSchema(
       name: type.name,
       description: type.description,
       types: type.getTypes().map(getTypeFromDef),
+      astNode: type.astNode,
       resolveType: type.resolveType,
     });
   }
@@ -397,6 +408,7 @@ export function extendSchema(
         deprecationReason: field.deprecationReason,
         type: extendFieldType(field.type),
         args: keyMap(field.args, arg => arg.name),
+        astNode: field.astNode,
         resolve: field.resolve,
       };
     });
@@ -419,6 +431,7 @@ export function extendSchema(
             type: buildOutputFieldType(field.type),
             args: buildInputValues(field.arguments),
             deprecationReason: getDeprecationReason(field),
+            astNode: field,
           };
         });
       });
@@ -456,6 +469,7 @@ export function extendSchema(
       description: getDescription(typeNode),
       interfaces: () => buildImplementedInterfaces(typeNode),
       fields: () => buildFieldMap(typeNode),
+      astNode: typeNode,
     });
   }
 
@@ -464,6 +478,7 @@ export function extendSchema(
       name: typeNode.name.value,
       description: getDescription(typeNode),
       fields: () => buildFieldMap(typeNode),
+      astNode: typeNode,
       resolveType: cannotExecuteExtendedSchema,
     });
   }
@@ -473,6 +488,7 @@ export function extendSchema(
       name: typeNode.name.value,
       description: getDescription(typeNode),
       types: typeNode.types.map(getObjectTypeFromAST),
+      astNode: typeNode,
       resolveType: cannotExecuteExtendedSchema,
     });
   }
@@ -481,6 +497,7 @@ export function extendSchema(
     return new GraphQLScalarType({
       name: typeNode.name.value,
       description: getDescription(typeNode),
+      astNode: typeNode,
       serialize: id => id,
       // Note: validation calls the parse functions to determine if a
       // literal value is correct. Returning null would cause use of custom
@@ -501,8 +518,10 @@ export function extendSchema(
         enumValue => ({
           description: getDescription(enumValue),
           deprecationReason: getDeprecationReason(enumValue),
+          astNode: enumValue,
         }),
       ),
+      astNode: typeNode,
     });
   }
 
@@ -511,6 +530,7 @@ export function extendSchema(
       name: typeNode.name.value,
       description: getDescription(typeNode),
       fields: () => buildInputValues(typeNode.fields),
+      astNode: typeNode,
     });
   }
 
@@ -524,6 +544,7 @@ export function extendSchema(
       ),
       args:
         directiveNode.arguments && buildInputValues(directiveNode.arguments),
+      astNode: directiveNode,
     });
   }
 
@@ -541,6 +562,7 @@ export function extendSchema(
         description: getDescription(field),
         args: buildInputValues(field.arguments),
         deprecationReason: getDeprecationReason(field),
+        astNode: field,
       })
     );
   }
@@ -554,7 +576,8 @@ export function extendSchema(
         return {
           type,
           description: getDescription(value),
-          defaultValue: valueFromAST(value.defaultValue, type)
+          defaultValue: valueFromAST(value.defaultValue, type),
+          astNode: value,
         };
       }
     );