Let extensions intercept schemas

src/extension.ts

@@ -10,7 +10,7 @@ export type GraphQLExtensionDeclaration = (
 export interface ExtensionAPI {
   logger: any;
   loaders: {
-    schema: Pick<LoadersRegistry, 'register'>;
+    schema: Pick<LoadersRegistry, 'register' | 'use'>;
     documents: Pick<LoadersRegistry, 'register'>;
   };
 }

src/helpers/utils.ts

@@ -30,3 +30,15 @@ export function isLegacyProjectConfig(
     typeof (config as IGraphQLProjectLegacy).excludes !== 'undefined'
   );
 }
+
+export type MiddlewareFn<T> = (input: T) => T;
+
+export function useMiddleware<T>(fns: Array<MiddlewareFn<T>>) {
+  return (input: T) => {
+    if (fns.length) {
+      return fns.reduce((obj, cb) => cb(obj), input);
+    }
+
+    return input;
+  };
+}

src/loaders.ts

@@ -4,18 +4,20 @@ import {
   LoadTypedefsOptions,
   loadDocuments,
   loadDocumentsSync,
-  loadSchema,
-  loadSchemaSync,
+  OPERATION_KINDS,
   loadTypedefs,
   loadTypedefsSync,
 } from '@graphql-toolkit/core';
-import {GraphQLSchema} from 'graphql';
+import {mergeTypeDefs} from '@graphql-toolkit/schema-merging';
+import {GraphQLSchema, DocumentNode, buildASTSchema} from 'graphql';
+import {MiddlewareFn, useMiddleware} from './helpers/utils';
 
 type Pointer = UnnormalizedTypeDefPointer | UnnormalizedTypeDefPointer[];
 type Options = Partial<LoadTypedefsOptions>;
 
 export class LoadersRegistry {
   private _loaders: Loader[] = [];
+  private _middlewares: MiddlewareFn<DocumentNode>[] = [];
   private readonly cwd: string;
 
   constructor({cwd}: {cwd: string}) {
@@ -28,6 +30,10 @@ export class LoadersRegistry {
     }
   }
 
+  use(middleware: MiddlewareFn<DocumentNode>): void {
+    this._middlewares.push(middleware);
+  }
+
   async loadTypeDefs(pointer: Pointer, options?: Options): Promise<Source[]> {
     return loadTypedefs(pointer, {
       loaders: this._loaders,
@@ -52,11 +58,21 @@ export class LoadersRegistry {
     pointer: Pointer,
     options?: Options,
   ): Promise<GraphQLSchema> {
-    return loadSchema(pointer, this.createOptions(options));
+    const sources = await loadTypedefs(pointer, {
+      filterKinds: OPERATION_KINDS,
+      ...this.createOptions(options),
+    });
+
+    return this.buildSchema(sources);
   }
 
   loadSchemaSync(pointer: Pointer, options?: Options): GraphQLSchema {
-    return loadSchemaSync(pointer, this.createOptions(options));
+    const sources = loadTypedefsSync(pointer, {
+      filterKinds: OPERATION_KINDS,
+      ...this.createOptions(options),
+    });
+
+    return this.buildSchema(sources);
   }
 
   private createOptions<T extends object>(options?: T) {
@@ -66,4 +82,11 @@ export class LoadersRegistry {
       ...options,
     };
   }
+
+  private buildSchema(sources: Source[]) {
+    const documents: DocumentNode[] = sources.map(source => source.document);
+    const document = mergeTypeDefs(documents);
+
+    return buildASTSchema(useMiddleware(this._middlewares)(document));
+  }
 }

test/loaders.spec.ts

@@ -0,0 +1,53 @@
+import {parse, DirectiveDefinitionNode} from 'graphql';
+
+jest.mock('@graphql-toolkit/core', () => {
+  return {
+    loadTypedefsSync() {
+      return [
+        {
+          document: parse(/* GraphQL */ `
+            type Query {
+              foo: String @cache
+            }
+          `),
+        },
+      ];
+    },
+  };
+});
+
+import {LoadersRegistry} from '../src/loaders';
+
+test('Middlewares', () => {
+  const registry = new LoadersRegistry({cwd: __dirname});
+
+  expect(() => {
+    registry.loadSchemaSync('anything');
+  }).toThrow(/cache/);
+
+  const cacheDirective: DirectiveDefinitionNode = {
+    kind: 'DirectiveDefinition',
+    name: {
+      kind: 'Name',
+      value: 'cache',
+    },
+    repeatable: false,
+    locations: [
+      {
+        kind: 'Name',
+        value: 'FIELD_DEFINITION',
+      },
+    ],
+  };
+
+  registry.use(doc => {
+    return {
+      ...doc,
+      definitions: [...doc.definitions, cacheDirective],
+    };
+  });
+
+  const schema = registry.loadSchemaSync('anything');
+
+  expect(schema.getDirective('cache')).toBeDefined();
+});