Refactor jit plugin

jest.config.js

@@ -1,5 +1,5 @@
 const { resolve } = require('path');
-const { pathsToModuleNameMapper } = require('ts-jest/utils');
+const { pathsToModuleNameMapper } = require('ts-jest');
 const CI = !!process.env.CI;
 
 const ROOT_DIR = __dirname;

packages/plugins/graphql-jit/src/index.ts

@@ -1,22 +1,29 @@
 /* eslint-disable no-console */
-import { Plugin, TypedExecutionArgs } from '@envelop/core';
-import { DocumentNode, Source, ExecutionArgs, ExecutionResult } from 'graphql';
-import { compileQuery, isCompiledQuery, CompilerOptions, CompiledQuery } from 'graphql-jit';
+import { EnvelopError, Plugin, PromiseOrValue } from '@envelop/core';
+import { DocumentNode, Source, ExecutionArgs, ExecutionResult, print, GraphQLSchema, getOperationAST } from 'graphql';
+import { compileQuery, CompilerOptions, CompiledQuery } from 'graphql-jit';
 import lru from 'tiny-lru';
 
 const DEFAULT_MAX = 1000;
 const DEFAULT_TTL = 3600000;
 
-type JITCacheEntry = {
-  query: CompiledQuery['query'];
-  subscribe?: CompiledQuery['subscribe'];
-};
+type CompilationResult = ReturnType<typeof compileQuery>;
 
 export interface JITCache {
-  get(key: string): JITCacheEntry | undefined;
-  set(key: string, value: JITCacheEntry): void;
+  get(key: string): CompilationResult | undefined;
+  set(key: string, value: CompilationResult): void;
+}
+
+function isSource(source: any): source is Source {
+  return source.body != null;
 }
 
+function isCompiledQuery(compilationResult: any): compilationResult is CompiledQuery<unknown, unknown> {
+  return compilationResult.query != null;
+}
+
+const MUST_DEFINE_OPERATION_NAME_ERR = 'Must provide operation name if query contains multiple operations.';
+
 export const useGraphQlJit = (
   compilerOptions: Partial<CompilerOptions> = {},
   pluginOptions: {
@@ -27,7 +34,7 @@ export const useGraphQlJit = (
     /**
      * Callback triggered in case of GraphQL Jit compilation error.
      */
-    onError?: (r: ExecutionResult) => void;
+    onError?: (r: ExecutionResult<any>) => void;
     /**
      * Custom cache instance
      */
@@ -36,67 +43,92 @@ export const useGraphQlJit = (
 ): Plugin => {
   const documentSourceMap = new WeakMap<DocumentNode, string>();
   const jitCache =
-    typeof pluginOptions.cache !== 'undefined' ? pluginOptions.cache : lru<JITCacheEntry>(DEFAULT_MAX, DEFAULT_TTL);
+    typeof pluginOptions.cache !== 'undefined' ? pluginOptions.cache : lru<CompilationResult>(DEFAULT_MAX, DEFAULT_TTL);
 
-  function getCacheEntry<T>(args: TypedExecutionArgs<T>): JITCacheEntry {
-    let cacheEntry: JITCacheEntry | undefined;
-    const documentSource = documentSourceMap.get(args.document);
+  function getCacheKey(document: DocumentNode, operationName?: string) {
+    if (!operationName) {
+      const operationAST = getOperationAST(document);
+      if (!operationAST) {
+        throw new EnvelopError(MUST_DEFINE_OPERATION_NAME_ERR);
+      }
+      operationName = operationAST.name?.value;
+    }
+    const documentSource = getDocumentSource(document);
+    return `${documentSource}::${operationName}`;
+  }
 
-    if (documentSource) {
-      cacheEntry = jitCache.get(documentSource);
+  function getDocumentSource(document: DocumentNode): string {
+    let documentSource = documentSourceMap.get(document);
+    if (!documentSource) {
+      documentSource = print(document);
+      documentSourceMap.set(document, documentSource);
     }
+    return documentSource;
+  }
 
-    if (!cacheEntry) {
-      const compilationResult = compileQuery(args.schema, args.document, args.operationName ?? undefined, compilerOptions);
+  function getCompilationResult(schema: GraphQLSchema, document: DocumentNode, operationName?: string) {
+    const cacheKey = getCacheKey(document, operationName);
 
-      if (!isCompiledQuery(compilationResult)) {
-        if (pluginOptions?.onError) {
-          pluginOptions.onError(compilationResult);
-        } else {
-          console.error(compilationResult);
-        }
-        cacheEntry = {
-          query: () => compilationResult,
-        };
-      } else {
-        cacheEntry = compilationResult;
-      }
+    let compiledQuery = jitCache.get(cacheKey);
 
-      if (documentSource) {
-        jitCache.set(documentSource, cacheEntry);
+    if (!compiledQuery) {
+      compiledQuery = compileQuery(schema, document, operationName, compilerOptions);
+      if (!isCompiledQuery(compiledQuery)) {
+        pluginOptions?.onError?.(compiledQuery);
       }
+      jitCache.set(cacheKey, compiledQuery);
     }
-    return cacheEntry;
+
+    return compiledQuery;
   }
 
   return {
     onParse({ params: { source } }) {
-      const key = source instanceof Source ? source.body : source;
+      const key = isSource(source) ? source.body : source;
 
       return ({ result }) => {
         if (!result || result instanceof Error) return;
 
         documentSourceMap.set(result, key);
       };
     },
-    async onExecute({ args, setExecuteFn }) {
+    onValidate({ params, setResult }) {
+      try {
+        const compilationResult = getCompilationResult(params.schema, params.documentAST);
+        if (!isCompiledQuery(compilationResult) && compilationResult?.errors != null) {
+          setResult(compilationResult.errors);
+        }
+      } catch (e: any) {
+        // Validate doesn't work in case of multiple operations
+        if (e.message !== MUST_DEFINE_OPERATION_NAME_ERR) {
+          throw e;
+        }
+      }
+    },
+    async onExecute({ args, setExecuteFn, setResultAndStopExecution }) {
       if (!pluginOptions.enableIf || (pluginOptions.enableIf && (await pluginOptions.enableIf(args)))) {
-        setExecuteFn(function jitExecutor() {
-          const cacheEntry = getCacheEntry(args);
-
-          return cacheEntry.query(args.rootValue, args.contextValue, args.variableValues);
-        });
+        const compilationResult = getCompilationResult(args.schema, args.document, args.operationName ?? undefined);
+        if (isCompiledQuery(compilationResult)) {
+          setExecuteFn(function jitExecute(args): PromiseOrValue<ExecutionResult<any, any>> {
+            return compilationResult.query(args.rootValue, args.contextValue, args.variableValues);
+          });
+        } else if (compilationResult != null) {
+          setResultAndStopExecution(compilationResult as ExecutionResult);
+        }
       }
     },
-    async onSubscribe({ args, setSubscribeFn }) {
+    async onSubscribe({ args, setSubscribeFn, setResultAndStopExecution }) {
       if (!pluginOptions.enableIf || (pluginOptions.enableIf && (await pluginOptions.enableIf(args)))) {
-        setSubscribeFn(async function jitSubscriber() {
-          const cacheEntry = getCacheEntry(args);
-
-          return cacheEntry.subscribe
-            ? (cacheEntry.subscribe(args.rootValue, args.contextValue, args.variableValues) as any)
-            : cacheEntry.query(args.rootValue, args.contextValue, args.variableValues);
-        });
+        const compilationResult = getCompilationResult(args.schema, args.document, args.operationName ?? undefined);
+        if (isCompiledQuery(compilationResult)) {
+          setSubscribeFn(function jitSubscribe(args: ExecutionArgs) {
+            return compilationResult.subscribe
+              ? compilationResult.subscribe(args.rootValue, args.contextValue, args.variableValues)
+              : compilationResult.query(args.rootValue, args.contextValue, args.variableValues);
+          } as any);
+        } else if (compilationResult != null) {
+          setResultAndStopExecution(compilationResult as ExecutionResult);
+        }
       }
     },
   };

packages/plugins/graphql-jit/test/graphql-jit.spec.ts

@@ -8,6 +8,7 @@ import { makeExecutableSchema } from '@graphql-tools/schema';
 import { execute, subscribe } from 'graphql';
 import { useGraphQlJit } from '../src';
 import lru from 'tiny-lru';
+import { GraphQLError } from 'graphql';
 
 describe('useGraphQlJit', () => {
   const schema = makeExecutableSchema({
@@ -52,7 +53,7 @@ describe('useGraphQlJit', () => {
 
     expect(onExecuteSpy).toHaveBeenCalledTimes(1);
     expect(onExecuteSpy.mock.calls[0][0].executeFn).not.toBe(execute);
-    expect(onExecuteSpy.mock.calls[0][0].executeFn.name).toBe('jitExecutor');
+    expect(onExecuteSpy.mock.calls[0][0].executeFn.name).toBe('jitExecute');
   });
 
   it('Should override subscribe function', async () => {
@@ -72,7 +73,7 @@ describe('useGraphQlJit', () => {
 
     expect(onSubscribeSpy).toHaveBeenCalledTimes(1);
     expect(onSubscribeSpy.mock.calls[0][0].subscribeFn).not.toBe(subscribe);
-    expect(onSubscribeSpy.mock.calls[0][0].subscribeFn.name).toBe('jitSubscriber');
+    expect(onSubscribeSpy.mock.calls[0][0].subscribeFn.name).toBe('jitSubscribe');
   });
 
   it('Should not override execute function when enableIf returns false', async () => {
@@ -97,7 +98,7 @@ describe('useGraphQlJit', () => {
 
     expect(onExecuteSpy).toHaveBeenCalledTimes(1);
     expect(onExecuteSpy.mock.calls[0][0].executeFn).toBe(execute);
-    expect(onExecuteSpy.mock.calls[0][0].executeFn.name).not.toBe('jitExecutor');
+    expect(onExecuteSpy.mock.calls[0][0].executeFn.name).not.toBe('jitExecute');
   });
 
   it('Should not override subscribe function when enableIf returns false', async () => {
@@ -122,7 +123,7 @@ describe('useGraphQlJit', () => {
 
     expect(onSubscribeSpy).toHaveBeenCalledTimes(1);
     expect(onSubscribeSpy.mock.calls[0][0].subscribeFn).toBe(subscribe);
-    expect(onSubscribeSpy.mock.calls[0][0].subscribeFn.name).not.toBe('jitSubscriber');
+    expect(onSubscribeSpy.mock.calls[0][0].subscribeFn.name).not.toBe('jitSubscribe');
   });
 
   it('Should execute correctly', async () => {
@@ -143,7 +144,7 @@ describe('useGraphQlJit', () => {
   });
 
   it('Should use the provided cache instance', async () => {
-    const cache = lru();
+    const cache = new Map();
     jest.spyOn(cache, 'set');
     jest.spyOn(cache, 'get');
 
@@ -160,7 +161,32 @@ describe('useGraphQlJit', () => {
     );
 
     await testInstance.execute(`query { test }`);
-    expect(cache.get).toHaveBeenCalled();
-    expect(cache.set).toHaveBeenCalled();
+    await testInstance.execute(`query { test }`);
+    expect(cache.get).toHaveBeenCalledTimes(4);
+    expect(cache.set).toHaveBeenCalledTimes(1);
+  });
+
+  it('Should throw validation errors if compilation fails', async () => {
+    const plugin = useGraphQlJit(
+      {},
+      {
+        cache: {
+          get() {
+            return {
+              errors: [new GraphQLError('Some random error')],
+            };
+          },
+          set() {},
+        },
+      }
+    );
+    jest.spyOn(plugin, 'onValidate');
+    jest.spyOn(plugin, 'onExecute');
+    const testInstance = createTestkit([plugin], schema);
+
+    const result = await testInstance.execute(`query Foo { test }`);
+    expect(result['errors']).toBeTruthy();
+    expect(plugin.onValidate).toHaveBeenCalled();
+    expect(plugin.onExecute).not.toHaveBeenCalled();
   });
 });