Update extendSchema to support new Directives (#455)

* Update extendSchema to support new Directives

This changes the extendSchema utility to accept and process new
Directive definitions. Existing directives cannot be extended or
replaced (should throw an error in such cases).

See extendSchema-test for examples.

* Update extendSchema to support new Directives

This changes the extendSchema utility to accept and process new
Directive definitions. Existing directives cannot be extended or
replaced (should throw an error in such cases).

See extendSchema-test for examples.

* Fix trailing whitespace lint error

* Improve logic in new directive extendSchema

Use an Array instead of an object for collecting new
directives detected within the schema extension.

Author: robzhu

src/utilities/__tests__/extendSchema-test.js

@@ -23,9 +23,9 @@ import {
   GraphQLEnumType,
   GraphQLNonNull,
   GraphQLList,
+  GraphQLScalarType,
 } from '../../type';
 
-
 // Test schema.
 const SomeInterfaceType = new GraphQLInterfaceType({
   name: 'SomeInterface',
@@ -718,6 +718,72 @@ type Subscription {
 `);
   });
 
+  it('may extend directives with new simple directive', () => {
+    const ast = parse(`
+      directive @neat on QUERY
+    `);
+
+    const extendedSchema = extendSchema(testSchema, ast);
+    const newDirective = extendedSchema.getDirective('neat');
+    expect(newDirective.name).to.equal('neat');
+    expect(newDirective.locations).to.contain('QUERY');
+  });
+
+  it('may extend directives with new complex directive', () => {
+    const ast = parse(`
+      directive @profile(enable: Boolean! tag: String) on QUERY | FIELD
+    `);
+
+    const extendedSchema = extendSchema(testSchema, ast);
+    const extendedDirective = extendedSchema.getDirective('profile');
+    expect(extendedDirective.locations).to.contain('QUERY');
+    expect(extendedDirective.locations).to.contain('FIELD');
+
+    const args = extendedDirective.args;
+    const arg0 = args[0];
+    const arg1 = args[1];
+
+    expect(args.length).to.equal(2);
+    expect(arg0.name).to.equal('enable');
+    expect(arg0.type).to.be.instanceof(GraphQLNonNull);
+    expect(arg0.type.ofType).to.be.instanceof(GraphQLScalarType);
+
+    expect(arg1.name).to.equal('tag');
+    expect(arg1.type).to.be.instanceof(GraphQLScalarType);
+  });
+
+  it('does not allow replacing a default directive', () => {
+    const ast = parse(`
+      directive @include(if: Boolean!) on FIELD | FRAGMENT_SPREAD
+    `);
+
+    expect(() =>
+      extendSchema(testSchema, ast)
+    ).to.throw(
+      'Directive "include" already exists in the schema. It cannot be ' +
+      'redefined.'
+    );
+  });
+
+  it('does not allow replacing a custom directive', () => {
+    const ast = parse(`
+      directive @meow(if: Boolean!) on FIELD | FRAGMENT_SPREAD
+    `);
+
+    const extendedSchema = extendSchema(testSchema, ast);
+
+    const replacementAst = parse(`
+      directive @meow(if: Boolean!) on FIELD | QUERY
+    `);
+
+    expect(() =>
+      extendSchema(extendedSchema, replacementAst)
+    ).to.throw(
+      'Directive "meow" already exists in the schema. It cannot be ' +
+      'redefined.'
+    );
+  });
+
   it('does not allow replacing an existing type', () => {
     const ast = parse(`
       type Bar {

src/utilities/extendSchema.js

@@ -28,6 +28,10 @@ import {
   isOutputType,
 } from '../type/definition';
 
+import {
+  GraphQLDirective,
+} from '../type/directives';
+
 import {
   __Schema,
   __Directive,
@@ -58,6 +62,7 @@ import {
   SCALAR_TYPE_DEFINITION,
   INPUT_OBJECT_TYPE_DEFINITION,
   TYPE_EXTENSION_DEFINITION,
+  DIRECTIVE_DEFINITION,
 } from '../language/kinds';
 
 import type {
@@ -67,6 +72,10 @@ import type {
   GraphQLOutputType,
 } from '../type/definition';
 
+import type {
+  DirectiveLocationEnum
+} from '../type/directives';
+
 import type {
   Document,
   InputValueDefinition,
@@ -79,6 +88,7 @@ import type {
   ScalarTypeDefinition,
   EnumTypeDefinition,
   InputObjectTypeDefinition,
+  DirectiveDefinition,
 } from '../language/ast';
 
 
@@ -112,6 +122,10 @@ export function extendSchema(
   const typeDefinitionMap = {};
   const typeExtensionsMap = {};
 
+  // New directives and types are separate because a directives and types can
+  // have the same name. For example, a type named "skip".
+  const directiveDefinitions : Array<DirectiveDefinition> = [];
+
   for (let i = 0; i < documentAST.definitions.length; i++) {
     const def = documentAST.definitions[i];
     switch (def.kind) {
@@ -159,13 +173,26 @@ export function extendSchema(
         }
         typeExtensionsMap[extendedTypeName] = extensions;
         break;
+      case DIRECTIVE_DEFINITION:
+        const directiveName = def.name.value;
+        const existingDirective = schema.getDirective(directiveName);
+        if (existingDirective) {
+          throw new GraphQLError(
+            `Directive "${directiveName}" already exists in the schema. It ` +
+            'cannot be redefined.',
+            [ def ]
+          );
+        }
+        directiveDefinitions.push(def);
+        break;
     }
   }
 
-  // If this document contains no new types, then return the same unmodified
-  // GraphQLSchema instance.
+  // If this document contains no new types, extensions, or directives then
+  // return the same unmodified GraphQLSchema instance.
   if (Object.keys(typeExtensionsMap).length === 0 &&
-      Object.keys(typeDefinitionMap).length === 0) {
+      Object.keys(typeDefinitionMap).length === 0 &&
+      directiveDefinitions.length === 0) {
     return schema;
   }
 
@@ -220,13 +247,22 @@ export function extendSchema(
     mutation: mutationType,
     subscription: subscriptionType,
     types,
-    // Copy directives.
-    directives: schema.getDirectives(),
+    directives: getMergedDirectives(),
   });
 
   // Below are functions used for producing this schema that have closed over
   // this scope and have access to the schema, cache, and newly defined types.
 
+  function getMergedDirectives(): Array<GraphQLDirective> {
+    const existingDirectives = schema.getDirectives();
+    invariant(existingDirectives, 'schema must have default directives');
+
+    const newDirectives = directiveDefinitions.map(directiveAST =>
+      getDirective(directiveAST)
+    );
+    return existingDirectives.concat(newDirectives);
+  }
+
   function getTypeFromDef<T: GraphQLNamedType>(typeDef: T): T {
     const type = _getNamedType(typeDef.name);
     invariant(type, 'Missing type from schema');
@@ -474,6 +510,18 @@ export function extendSchema(
     });
   }
 
+  function getDirective(
+    directiveAST: DirectiveDefinition
+  ): GraphQLDirective {
+    return new GraphQLDirective({
+      name: directiveAST.name.value,
+      locations: directiveAST.locations.map(
+        node => ((node.value: any): DirectiveLocationEnum)
+      ),
+      args: directiveAST.arguments && buildInputValues(directiveAST.arguments),
+    });
+  }
+
   function buildImplementedInterfaces(typeAST: ObjectTypeDefinition) {
     return typeAST.interfaces &&
       typeAST.interfaces.map(getInterfaceTypeFromAST);