Repository: mercurius-js/mercurius
Pull request #462 (2021-04-21 17:13:16)

Produce a friendly error message in case of a missing @key


Patch:
diff --git a/docs/hooks.md b/docs/hooks.md
index f4eea83..4de8e2d 100644
--- a/docs/hooks.md
+++ b/docs/hooks.md
@@ -230,3 +230,30 @@ fastify.graphql.addHook('preSubscriptionParsing', async (schema, source, context
   throw new Error('Some error')
 })
 ```
+
+## GraphQL Application lifecycle Hooks
+
+There is one hook that you can use in a GraphQL application.
+
+When registering hooks, you must wait for Mercurius to be registered in Fastify.
+
+```js
+await fastify.ready()
+```
+
+### onGatewayReplaceSchema
+
+When the Gateway service obtains new versions of federated schemas within a defined polling interval, the `onGatewayReplaceSchema` hook will be triggered every time a new schema is built. It is called just before the old schema is replaced with the new one.
+
+This hook will only be triggered in gateway mode. It has the following parameters:
+
+ - `instance` - The gateway server `FastifyInstance` (this contains the old schema).
+ - `schema` - The new schema that has been built from the gateway refresh.
+
+```js
+fastify.graphql.addHook('onGatewayReplaceSchema', async (instance, schema) => {
+  await someSchemaTraversalFn()
+})
+```
+
+If this hook throws, the error will be caught and logged using the `FastifyInstance` logger. Subsequent `onGatewayReplaceSchema` hooks that are registered will not be run for this interval.
diff --git a/examples/hooks-gateway.js b/examples/hooks-gateway.js
index 1c1a3fb..55fd22c 100644
--- a/examples/hooks-gateway.js
+++ b/examples/hooks-gateway.js
@@ -216,6 +216,10 @@ async function start () {
     console.log('onResolution called')
   })
 
+  gateway.graphql.addHook('onGatewayReplaceSchema', async (instance, schema) => {
+    console.log('onGatewayReplaceSchema called')
+  })
+
   await gateway.listen(4000)
 }
 
diff --git a/index.d.ts b/index.d.ts
index 50a5c4b..7323efc 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -203,6 +203,23 @@ export interface onSubscriptionEndHookHandler<TContext = MercuriusContext> {
   ): Promise<void>;
 }
 
+// ----------------------------
+// Application Lifecycle hooks
+// ----------------------------
+
+/**
+ * `onGatewayReplaceSchema` is an application lifeycle hook. When the Gateway service obtains new versions of federated schemas within a defined polling interval, the `onGatewayReplaceSchema` hook will be triggered every time a new schema is built. It is called just before the old schema is replaced with the new one.
+ * This hook will only be triggered in gateway mode. It has the following parameters:
+ *  - `instance` - The gateway server `FastifyInstance` (this contains the old schema).
+ *  - `schema` - The new schema that has been built from the gateway refresh.
+ */
+export interface onGatewayReplaceSchemaHookHandler {
+  (
+    instance: FastifyInstance,
+    schema: GraphQLSchema
+  ): Promise<void>;
+}
+
 interface ServiceConfig {
   setSchema: (schema: string) => ServiceConfig;
 }
@@ -328,6 +345,16 @@ interface MercuriusPlugin {
    * This hook will only be triggered when subscriptions are enabled.
    */
   addHook<TContext = MercuriusContext>(name: 'onSubscriptionEnd', hook: onSubscriptionEndHookHandler<TContext>): void;
+
+  // Application lifecycle addHooks
+
+  /**
+   * `onGatewayReplaceSchema` is an application lifeycle hook. When the Gateway service obtains new versions of federated schemas within a defined polling interval, the `onGatewayReplaceSchema` hook will be triggered every time a new schema is built. It is called just before the old schema is replaced with the new one.
+   * This hook will only be triggered in gateway mode. It has the following parameters:
+   *  - `instance` - The gateway server `FastifyInstance` (this contains the old schema).
+   *  - `schema` - The new schema that has been built from the gateway refresh.
+   */
+  addHook(name: 'onGatewayReplaceSchema', hook: onGatewayReplaceSchemaHookHandler): void;
 }
 
 interface QueryRequest {
diff --git a/index.js b/index.js
index d0c420f..9b0f70b 100644
--- a/index.js
+++ b/index.js
@@ -39,9 +39,9 @@ const {
   MER_ERR_METHOD_NOT_ALLOWED,
   MER_ERR_INVALID_METHOD
 } = require('./lib/errors')
-const { Hooks, assignLifeCycleHooksToContext } = require('./lib/hooks')
+const { Hooks, assignLifeCycleHooksToContext, assignApplicationLifecycleHooksToContext } = require('./lib/hooks')
 const { kLoaders, kFactory, kHooks } = require('./lib/symbols')
-const { preParsingHandler, preValidationHandler, preExecutionHandler, onResolutionHandler } = require('./lib/handlers')
+const { preParsingHandler, preValidationHandler, preExecutionHandler, onResolutionHandler, onGatewayReplaceSchemaHandler } = require('./lib/handlers')
 
 function buildCache (opts) {
   if (Object.prototype.hasOwnProperty.call(opts, 'cache')) {
@@ -179,9 +179,18 @@ const plugin = fp(async function (app, opts) {
     if (gateway.pollingInterval !== undefined) {
       if (typeof gateway.pollingInterval === 'number') {
         gatewayInterval = setInterval(async () => {
-          const schema = await gateway.refresh()
-          if (schema !== null) {
-            fastifyGraphQl.replaceSchema(schema)
+          try {
+            const context = assignApplicationLifecycleHooksToContext(fastifyGraphQl[kHooks], {})
+            const schema = await gateway.refresh()
+            if (schema !== null) {
+              // Trigger onGatewayReplaceSchema hook
+              if (context.onGatewayReplaceSchema !== null) {
+                await onGatewayReplaceSchemaHandler(context, { instance: app, schema })
+              }
+              fastifyGraphQl.replaceSchema(schema)
+            }
+          } catch (error) {
+            app.log.error(error)
           }
         }, gateway.pollingInterval)
       } else {
@@ -486,19 +495,10 @@ const plugin = fp(async function (app, opts) {
       }
     }
 
-    // minJit is 0 by default
-    if (cached && cached.count++ === minJit) {
-      cached.jit = compileQuery(fastifyGraphQl.schema, document, operationName)
-    }
-
-    if (cached && cached.jit !== null) {
-      const execution = await cached.jit.query(root, context, variables || {})
-
-      return maybeFormatErrors(execution, context)
-    }
+    const shouldCompileJit = cached && cached.count++ === minJit
 
     // Validate variables
-    if (variables !== undefined) {
+    if (variables !== undefined && !shouldCompileJit) {
       const executionContext = buildExecutionContext(fastifyGraphQl.schema, document, root, context, variables, operationName)
       if (Array.isArray(executionContext)) {
         const err = new MER_ERR_GQL_VALIDATION()
@@ -513,6 +513,17 @@ const plugin = fp(async function (app, opts) {
       ({ modifiedDocument } = await preExecutionHandler({ schema: fastifyGraphQl.schema, document, context }))
     }
 
+    // minJit is 0 by default
+    if (shouldCompileJit) {
+      cached.jit = compileQuery(fastifyGraphQl.schema, modifiedDocument || document, operationName)
+    }
+
+    if (cached && cached.jit !== null) {
+      const execution = await cached.jit.query(root, context, variables || {})
+
+      return maybeFormatErrors(execution, context)
+    }
+
     const execution = await execute(
       fastifyGraphQl.schema,
       modifiedDocument || document,
diff --git a/lib/errors.js b/lib/errors.js
index 57f371c..c83b789 100644
--- a/lib/errors.js
+++ b/lib/errors.js
@@ -169,6 +169,10 @@ const errors = {
     'MER_ERR_GQL_GATEWAY_INIT',
     'Gateway schema init issues'
   ),
+  MER_ERR_GQL_GATEWAY_MISSING_KEY_DIRECTIVE: createError(
+    'MER_ERR_GQL_GATEWAY_MISSING_KEY_DIRECTIVE',
+    'Missing @key directive in %s type'
+  ),
   /**
    * Persisted query errors
    */
diff --git a/lib/gateway/make-resolver.js b/lib/gateway/make-resolver.js
index 062dbc3..58f9ea3 100644
--- a/lib/gateway/make-resolver.js
+++ b/lib/gateway/make-resolver.js
@@ -8,6 +8,10 @@ const {
 } = require('graphql')
 const { preGatewayExecutionHandler, preGatewaySubscriptionExecutionHandler } = require('../handlers')
 
+const {
+  MER_ERR_GQL_GATEWAY_MISSING_KEY_DIRECTIVE
+} = require('../errors')
+
 const kEntityResolvers = Symbol('mercurius.entity-resolvers')
 
 function getFieldType (schema, type, fieldName) {
@@ -489,6 +493,10 @@ function makeResolver ({ service, createOperation, transformData, isQuery, isRef
 function removeNonIdProperties (obj, type) {
   const keyDirective = type.astNode.directives.find(d => d.name.value === 'key')
 
+  if (!keyDirective) {
+    throw new MER_ERR_GQL_GATEWAY_MISSING_KEY_DIRECTIVE(type.name)
+  }
+
   const idFields = keyDirective.arguments[0].value.value.split(' ')
 
   const result = {
diff --git a/lib/handlers.js b/lib/handlers.js
index cf8648d..1c8d2b4 100644
--- a/lib/handlers.js
+++ b/lib/handlers.js
@@ -1,9 +1,17 @@
 'use strict'
 
-const { hooksRunner, preExecutionHooksRunner, preGatewayExecutionHooksRunner, gatewayHookRunner, hookRunner, preParsingHookRunner, onResolutionHookRunner, onEndHookRunner } = require('./hooks')
+const { hooksRunner, preExecutionHooksRunner, preGatewayExecutionHooksRunner, gatewayHookRunner, hookRunner, preParsingHookRunner, onResolutionHookRunner, onEndHookRunner, onGatewayReplaceSchemaHookRunner } = require('./hooks')
 const { addErrorsToContext } = require('./errors')
 const { print } = require('graphql')
 
+async function onGatewayReplaceSchemaHandler (context, data) {
+  await hooksRunner(
+    context.onGatewayReplaceSchema,
+    onGatewayReplaceSchemaHookRunner,
+    data
+  )
+}
+
 async function preParsingHandler (request) {
   await hooksRunner(
     request.context.preParsing,
@@ -106,5 +114,6 @@ module.exports = {
   preSubscriptionExecutionHandler,
   preGatewaySubscriptionExecutionHandler,
   onSubscriptionResolutionHandler,
-  onSubscriptionEndHandler
+  onSubscriptionEndHandler,
+  onGatewayReplaceSchemaHandler
 }
diff --git a/lib/hooks.js b/lib/hooks.js
index eb85a6d..3060489 100644
--- a/lib/hooks.js
+++ b/lib/hooks.js
@@ -1,5 +1,8 @@
 'use strict'
 
+const applicationHooks = [
+  'onGatewayReplaceSchema'
+]
 const lifecycleHooks = [
   'preParsing',
   'preValidation',
@@ -12,6 +15,7 @@ const lifecycleHooks = [
   'onSubscriptionResolution',
   'onSubscriptionEnd'
 ]
+const supportedHooks = lifecycleHooks.concat(applicationHooks)
 const { MER_ERR_HOOK_INVALID_TYPE, MER_ERR_HOOK_INVALID_HANDLER, MER_ERR_HOOK_UNSUPPORTED_HOOK } = require('./errors')
 
 function Hooks () {
@@ -25,12 +29,13 @@ function Hooks () {
   this.preGatewaySubscriptionExecution = []
   this.onSubscriptionResolution = []
   this.onSubscriptionEnd = []
+  this.onGatewayReplaceSchema = []
 }
 
 Hooks.prototype.validate = function (hook, fn) {
   if (typeof hook !== 'string') throw new MER_ERR_HOOK_INVALID_TYPE()
   if (typeof fn !== 'function') throw new MER_ERR_HOOK_INVALID_HANDLER()
-  if (lifecycleHooks.indexOf(hook) === -1) {
+  if (supportedHooks.indexOf(hook) === -1) {
     throw new MER_ERR_HOOK_UNSUPPORTED_HOOK(hook)
   }
 }
@@ -40,6 +45,14 @@ Hooks.prototype.add = function (hook, fn) {
   this[hook].push(fn)
 }
 
+function assignApplicationLifecycleHooksToContext (hooks, context) {
+  const contextHooks = {
+    onGatewayReplaceSchema: null
+  }
+  if (hooks.onGatewayReplaceSchema.length > 0) contextHooks.onGatewayReplaceSchema = hooks.onGatewayReplaceSchema.slice()
+  return Object.assign(contextHooks, context)
+}
+
 function assignLifeCycleHooksToContext (hooks, context) {
   const contextHooks = {
     preParsing: null,
@@ -132,9 +145,14 @@ function onEndHookRunner (fn, request) {
   return fn(request.context)
 }
 
+function onGatewayReplaceSchemaHookRunner (fn, data) {
+  return fn(data.instance, data.schema)
+}
+
 module.exports = {
   Hooks,
   assignLifeCycleHooksToContext,
+  assignApplicationLifecycleHooksToContext,
   hooksRunner,
   preExecutionHooksRunner,
   hookRunner,
@@ -143,5 +161,5 @@ module.exports = {
   preParsingHookRunner,
   onResolutionHookRunner,
   onEndHookRunner,
-  lifecycleHooks
+  onGatewayReplaceSchemaHookRunner
 }


Test patch:
diff --git a/test/gateway/application-hooks.js b/test/gateway/application-hooks.js
new file mode 100644
index 0000000..c18f5f7
--- /dev/null
+++ b/test/gateway/application-hooks.js
@@ -0,0 +1,203 @@
+'use strict'
+
+const { test } = require('tap')
+const FakeTimers = require('@sinonjs/fake-timers')
+const { GraphQLSchema } = require('graphql')
+const { promisify } = require('util')
+const Fastify = require('fastify')
+const GQL = require('../..')
+const buildFederationSchema = require('../../lib/federation')
+
+const immediate = promisify(setImmediate)
+
+// ----------------------
+// onGatewayReplaceSchema
+// ----------------------
+test('onGatewayReplaceSchema - polling interval with a new schema should trigger onGatewayReplaceSchema hook', async (t) => {
+  t.plan(2)
+
+  const clock = FakeTimers.install({
+    shouldAdvanceTime: true,
+    advanceTimeDelta: 40
+  })
+  t.teardown(() => clock.uninstall())
+
+  const resolvers = {
+    Query: {
+      me: (root, args, context, info) => user
+    },
+    User: {
+      __resolveReference: (user, args, context, info) => user
+    }
+  }
+
+  const user = {
+    id: 'u1',
+    name: 'John',
+    lastName: 'Doe'
+  }
+
+  const userService = Fastify()
+  const gateway = Fastify()
+  t.teardown(async () => {
+    await gateway.close()
+    await userService.close()
+  })
+
+  userService.register(GQL, {
+    schema: `
+      extend type Query {
+        me: User
+      }
+
+      type User @key(fields: "id") {
+        id: ID!
+        name: String!
+      }
+    `,
+    resolvers: resolvers,
+    federationMetadata: true
+  })
+
+  await userService.listen(0)
+
+  const userServicePort = userService.server.address().port
+
+  await gateway.register(GQL, {
+    gateway: {
+      services: [
+        {
+          name: 'user',
+          url: `http://localhost:${userServicePort}/graphql`
+        }
+      ],
+      pollingInterval: 2000
+    }
+  })
+
+  gateway.graphql.addHook('onGatewayReplaceSchema', async (instance, schema) => {
+    t.type(instance, 'object')
+    t.type(schema, GraphQLSchema)
+  })
+
+  userService.graphql.replaceSchema(
+    buildFederationSchema(`
+      extend type Query {
+        me: User
+      }
+
+      type User @key(fields: "id") {
+        id: ID!
+        name: String!
+        lastName: String!
+      }
+    `)
+  )
+  userService.graphql.defineResolvers(resolvers)
+
+  await clock.tickAsync(2000)
+
+  // We need the event loop to actually spin twice to
+  // be able to propagate the change
+  await immediate()
+  await immediate()
+})
+
+test('onGatewayReplaceSchema - should log an error should any errors occur in the hook', async (t) => {
+  t.plan(2)
+
+  const clock = FakeTimers.install({
+    shouldAdvanceTime: true,
+    advanceTimeDelta: 40
+  })
+  t.teardown(() => clock.uninstall())
+
+  const resolvers = {
+    Query: {
+      me: (root, args, context, info) => user
+    },
+    User: {
+      __resolveReference: (user, args, context, info) => user
+    }
+  }
+
+  const user = {
+    id: 'u1',
+    name: 'John',
+    lastName: 'Doe'
+  }
+
+  const userService = Fastify()
+  const gateway = Fastify()
+  t.teardown(async () => {
+    await gateway.close()
+    await userService.close()
+  })
+
+  userService.register(GQL, {
+    schema: `
+      extend type Query {
+        me: User
+      }
+
+      type User @key(fields: "id") {
+        id: ID!
+        name: String!
+      }
+    `,
+    resolvers: resolvers,
+    federationMetadata: true
+  })
+
+  await userService.listen(0)
+
+  const userServicePort = userService.server.address().port
+
+  await gateway.register(GQL, {
+    gateway: {
+      services: [
+        {
+          name: 'user',
+          url: `http://localhost:${userServicePort}/graphql`
+        }
+      ],
+      pollingInterval: 2000
+    }
+  })
+
+  gateway.graphql.addHook('onGatewayReplaceSchema', async (instance, schema) => {
+    t.ok('trigger error')
+    throw new Error('kaboom')
+  })
+
+  gateway.graphql.addHook('onGatewayReplaceSchema', async (instance, schema) => {
+    t.fail('should not be called')
+  })
+
+  // Override gateway error logger
+  gateway.log.error = (error) => {
+    t.same(error, new Error('kaboom'))
+  }
+
+  userService.graphql.replaceSchema(
+    buildFederationSchema(`
+      extend type Query {
+        me: User
+      }
+
+      type User @key(fields: "id") {
+        id: ID!
+        name: String!
+        lastName: String!
+      }
+    `)
+  )
+  userService.graphql.defineResolvers(resolvers)
+
+  await clock.tickAsync(2000)
+
+  // We need the event loop to actually spin twice to
+  // be able to propagate the change
+  await immediate()
+  await immediate()
+})
diff --git a/test/gateway/extends-directive.js b/test/gateway/extends-directive.js
index 5fc316a..9a73b75 100644
--- a/test/gateway/extends-directive.js
+++ b/test/gateway/extends-directive.js
@@ -125,8 +125,6 @@ test('gateway handles @extends directive correctly', async (t) => {
     }
   })
 
-  await gateway.listen(0)
-
   const query = `
     query {
       me {
@@ -256,8 +254,6 @@ test('gateway passes field arguments through to types labeled by @extends direct
     }
   })
 
-  await gateway.listen(0)
-
   const query = `
     query {
       me {
@@ -475,3 +471,130 @@ test('gateway distributes query correctly to services when querying with inline
     }
   })
 })
+
+test('gateway handles missing @key', async (t) => {
+  // This service is missing a @key
+  const [userService, userServicePort] = await createService(t, `
+    type Query @extends {
+      me: User
+    }
+
+    type User {
+      id: ID!
+      name: String!
+    }
+  `, {
+    Query: {
+      me: (root, args, context, info) => {
+        return users.u1
+      }
+    },
+    User: {
+      __resolveReference: (user, args, context, info) => {
+        return users[user.id]
+      }
+    }
+  })
+
+  const [postService, postServicePort] = await createService(t, `
+    type Post @key(fields: "pid") {
+      pid: ID!
+      title: String
+      content: String
+      author: User @requires(fields: "title")
+    }
+
+    extend type Query {
+      topPosts(count: Int): [Post]
+    }
+
+    type User @key(fields: "id") @extends {
+      id: ID! @external
+      posts: [Post]
+      numberOfPosts: Int
+    }
+  `, {
+    Post: {
+      __resolveReference: (post, args, context, info) => {
+        return posts[post.pid]
+      },
+      author: (post, args, context, info) => {
+        return {
+          __typename: 'User',
+          id: post.authorId
+        }
+      }
+    },
+    User: {
+      posts: (user, args, context, info) => {
+        return Object.values(posts).filter(p => p.authorId === user.id)
+      },
+      numberOfPosts: (user) => {
+        return Object.values(posts).filter(p => p.authorId === user.id).length
+      }
+    },
+    Query: {
+      topPosts: (root, { count = 2 }) => Object.values(posts).slice(0, count)
+    }
+  })
+
+  const gateway = Fastify()
+  t.teardown(async () => {
+    await gateway.close()
+    await postService.close()
+    await userService.close()
+  })
+  gateway.register(GQL, {
+    gateway: {
+      services: [{
+        name: 'user',
+        url: `http://localhost:${userServicePort}/graphql`
+      }, {
+        name: 'post',
+        url: `http://localhost:${postServicePort}/graphql`
+      }]
+    }
+  })
+
+  const query = `
+    query {
+      me {
+        id
+        name
+        numberOfPosts
+      }
+    }
+  `
+
+  const res = await gateway.inject({
+    method: 'POST',
+    headers: {
+      'content-type': 'application/json'
+    },
+    url: '/graphql',
+    body: JSON.stringify({
+      query
+    })
+  })
+
+  t.same(JSON.parse(res.body), {
+    data: {
+      me: {
+        id: 'u1',
+        name: 'John',
+        numberOfPosts: null
+      }
+    },
+    errors: [{
+      message: 'Missing @key directive in User type',
+      locations: [{
+        line: 6,
+        column: 9
+      }],
+      path: [
+        'me',
+        'numberOfPosts'
+      ]
+    }]
+  })
+})
diff --git a/test/gateway/hooks.js b/test/gateway/hooks.js
index d395bee..57f31f1 100644
--- a/test/gateway/hooks.js
+++ b/test/gateway/hooks.js
@@ -75,7 +75,7 @@ const query = `
   }
 `
 
-async function createTestGatewayServer (t) {
+async function createTestGatewayServer (t, opts = {}) {
   // User service
   const userServiceSchema = `
   type Query @extends {
@@ -145,6 +145,7 @@ async function createTestGatewayServer (t) {
     await postService.close()
   })
   gateway.register(GQL, {
+    ...opts,
     gateway: {
       services: [{
         name: 'user',
@@ -283,6 +284,136 @@ test('gateway - hooks validation should handle invalid hook handlers', async (t)
   }
 })
 
+test('gateway - hooks should trigger when JIT is enabled', async (t) => {
+  t.plan(60)
+  const app = await createTestGatewayServer(t, { jit: 1 })
+
+  app.graphql.addHook('preParsing', async function (schema, source, context) {
+    await immediate()
+    t.type(schema, GraphQLSchema)
+    t.equal(source, query)
+    t.type(context, 'object')
+    t.ok('preParsing called')
+  })
+
+  // preValidation is not triggered a second time
+  app.graphql.addHook('preValidation', async function (schema, document, context) {
+    await immediate()
+    t.type(schema, GraphQLSchema)
+    t.same(document, parse(query))
+    t.type(context, 'object')
+    t.ok('preValidation called')
+  })
+
+  app.graphql.addHook('preExecution', async function (schema, document, context) {
+    await immediate()
+    t.type(schema, GraphQLSchema)
+    t.same(document, parse(query))
+    t.type(context, 'object')
+    t.ok('preExecution called')
+  })
+
+  // Execution events:
+  //  - once for user service query
+  //  - once for post service query
+  //  - once for reference type topPosts on User
+  //  - once for reference type author on Post
+  app.graphql.addHook('preGatewayExecution', async function (schema, document, context) {
+    await immediate()
+    t.type(schema, GraphQLSchema)
+    t.type(document, 'object')
+    t.type(context, 'object')
+    t.ok('preGatewayExecution called')
+  })
+
+  app.graphql.addHook('onResolution', async function (execution, context) {
+    await immediate()
+    t.type(execution, 'object')
+    t.type(context, 'object')
+    t.ok('onResolution called')
+  })
+
+  {
+    const res = await app.inject({
+      method: 'POST',
+      headers: { 'content-type': 'application/json' },
+      url: '/graphql',
+      body: JSON.stringify({ query })
+    })
+
+    t.same(JSON.parse(res.body), {
+      data: {
+        me: {
+          id: 'u1',
+          name: 'John',
+          topPosts: [
+            {
+              pid: 'p1',
+              author: {
+                id: 'u1'
+              }
+            },
+            {
+              pid: 'p3',
+              author: {
+                id: 'u1'
+              }
+            }
+          ]
+        },
+        topPosts: [
+          {
+            pid: 'p1'
+          },
+          {
+            pid: 'p2'
+          }
+        ]
+      }
+    })
+  }
+
+  {
+    const res = await app.inject({
+      method: 'POST',
+      headers: { 'content-type': 'application/json' },
+      url: '/graphql',
+      body: JSON.stringify({ query })
+    })
+
+    t.same(JSON.parse(res.body), {
+      data: {
+        me: {
+          id: 'u1',
+          name: 'John',
+          topPosts: [
+            {
+              pid: 'p1',
+              author: {
+                id: 'u1'
+              }
+            },
+            {
+              pid: 'p3',
+              author: {
+                id: 'u1'
+              }
+            }
+          ]
+        },
+        topPosts: [
+          {
+            pid: 'p1'
+          },
+          {
+            pid: 'p2'
+          }
+        ]
+      }
+    })
+  }
+})
+
 // --------------------
 // preParsing
 // --------------------
diff --git a/test/gateway/rewrite-headers.js b/test/gateway/rewrite-headers.js
index 7dd479f..addfad7 100644
--- a/test/gateway/rewrite-headers.js
+++ b/test/gateway/rewrite-headers.js
@@ -39,14 +39,20 @@ async function createUserService ({ hooks } = {}) {
     }
   }
 
-  return await createTestService(schema, resolvers, hooks)
+  return createTestService(schema, resolvers, hooks)
 }
 
 test('gateway - service rewriteHeaders', async (t) => {
   t.test('rewriteHeaders is called as expected', async (t) => {
     t.plan(5)
 
+    const [users, usersPort] = await createUserService()
+
     const gateway = Fastify()
+    t.teardown(async () => {
+      await gateway.close()
+      await users.close()
+    })
 
     const rewriteHeaders = (headers, context = 'not-passed') => {
       t.ok(headers != null, 'Headers is never undefined/null')
@@ -56,15 +62,9 @@ test('gateway - service rewriteHeaders', async (t) => {
       t.ok(context === 'not-passed' || context.app === gateway)
     }
 
-    const [users, usersPort] = await createUserService()
     const url = `http://localhost:${usersPort}/graphql`
     gateway.register(GQL, { gateway: { services: [{ name: 'user', url, rewriteHeaders }] } })
 
-    t.teardown(async () => {
-      await users.close()
-      await gateway.close()
-    })
-
     const res = await gateway.inject({
       method: 'POST',
       headers: { 'content-type': 'application/json' },
@@ -85,16 +85,17 @@ test('gateway - service rewriteHeaders', async (t) => {
     }
 
     const [users, usersPort] = await createUserService({ hooks: { onRequest } })
-    const url = `http://localhost:${usersPort}/graphql`
-    const gateway = Fastify()
-    const rewriteHeaders = () => ({ 'x-custom': custom })
-    gateway.register(GQL, { gateway: { services: [{ name: 'user', url, rewriteHeaders }] } })
 
+    const gateway = Fastify()
     t.teardown(async () => {
-      await users.close()
       await gateway.close()
+      await users.close()
     })
 
+    const rewriteHeaders = () => ({ 'x-custom': custom })
+    const url = `http://localhost:${usersPort}/graphql`
+    gateway.register(GQL, { gateway: { services: [{ name: 'user', url, rewriteHeaders }] } })
+
     const res = await gateway.inject({
       method: 'POST',
       headers: { 'content-type': 'application/json' },
diff --git a/test/hooks.js b/test/hooks.js
index 20e9bed..47f4ec6 100644
--- a/test/hooks.js
+++ b/test/hooks.js
@@ -27,11 +27,11 @@ const resolvers = {
 
 const query = '{ add(x: 2, y: 2) }'
 
-function createTestServer (t, customResolvers = resolvers) {
+function createTestServer (t, customResolvers = resolvers, opts = {}) {
   const app = Fastify()
   t.teardown(() => app.close())
 
-  app.register(GQL, { schema, resolvers: customResolvers })
+  app.register(GQL, { schema, resolvers: customResolvers, ...opts })
 
   return app
 }
@@ -161,6 +161,77 @@ test('hooks validation should handle invalid hook handlers', async t => {
   }
 })
 
+test('hooks should trigger when JIT is enabled', async t => {
+  t.plan(28)
+  const app = await createTestServer(t, resolvers, { jit: 1 })
+
+  app.graphql.addHook('preParsing', async function (schema, source, context) {
+    await immediate()
+    t.type(schema, GraphQLSchema)
+    t.equal(source, query)
+    t.type(context, 'object')
+    t.ok('preParsing called')
+  })
+
+  // preValidation is not triggered a second time
+  app.graphql.addHook('preValidation', async function (schema, document, context) {
+    await immediate()
+    t.type(schema, GraphQLSchema)
+    t.same(document, parse(query))
+    t.type(context, 'object')
+    t.ok('preValidation called')
+  })
+
+  app.graphql.addHook('preExecution', async function (schema, document, context) {
+    await immediate()
+    t.type(schema, GraphQLSchema)
+    t.same(document, parse(query))
+    t.type(context, 'object')
+    t.ok('preExecution called')
+  })
+
+  app.graphql.addHook('preGatewayExecution', async function (schema, document, context) {
+    t.fail('this should not be called')
+  })
+
+  app.graphql.addHook('onResolution', async function (execution, context) {
+    await immediate()
+    t.type(execution, 'object')
+    t.type(context, 'object')
+    t.ok('onResolution called')
+  })
+
+  {
+    const res = await app.inject({
+      method: 'POST',
+      headers: { 'content-type': 'application/json' },
+      url: '/graphql',
+      body: JSON.stringify({ query })
+    })
+
+    t.same(JSON.parse(res.body), {
+      data: {
+        add: 4
+      }
+    })
+  }
+
+  {
+    const res = await app.inject({
+      method: 'POST',
+      headers: { 'content-type': 'application/json' },
+      url: '/graphql',
+      body: JSON.stringify({ query })
+    })
+
+    t.same(JSON.parse(res.body), {
+      data: {
+        add: 4
+      }
+    })
+  }
+})
+
 // -----------
 // preParsing
 // -----------
diff --git a/test/types/index.ts b/test/types/index.ts
index 034b15e..2f3efac 100644
--- a/test/types/index.ts
+++ b/test/types/index.ts
@@ -1,13 +1,14 @@
+import { expectAssignable } from 'tsd'
 /* eslint-disable no-unused-expressions */
 import { EventEmitter } from 'events'
 // eslint-disable-next-line no-unused-vars
-import Fastify, { FastifyReply, FastifyRequest } from 'fastify'
+import Fastify, { FastifyReply, FastifyRequest, FastifyInstance } from 'fastify'
 // eslint-disable-next-line no-unused-vars
 import { Readable } from 'stream'
 // eslint-disable-next-line no-unused-vars
-import mercurius, { MercuriusOptions, IResolvers } from '../..'
+import mercurius, { MercuriusOptions, IResolvers, MercuriusContext, MercuriusServiceMetadata } from '../..'
 // eslint-disable-next-line no-unused-vars
-import { ValidationContext, ValidationRule } from 'graphql'
+import { DocumentNode, ExecutionResult, GraphQLSchema, ValidationContext, ValidationRule } from 'graphql'
 import { makeExecutableSchema } from '@graphql-tools/schema'
 import { mapSchema } from '@graphql-tools/utils'
 import mq from 'mqemitter'
@@ -467,15 +468,21 @@ query hello {
 `)
 
 app.graphql.addHook('preParsing', async function (schema, source, context) {
-  console.log('preParsing called')
+  expectAssignable<GraphQLSchema>(schema)
+  expectAssignable<string>(source)
+  expectAssignable<MercuriusContext>(context)
 })
 
 app.graphql.addHook('preValidation', async function (schema, document, context) {
-  console.log('preValidation called')
+  expectAssignable<GraphQLSchema>(schema)
+  expectAssignable<DocumentNode>(document)
+  expectAssignable<MercuriusContext>(context)
 })
 
 app.graphql.addHook('preExecution', async function (schema, document, context) {
-  console.log('preExecution called')
+  expectAssignable<GraphQLSchema>(schema)
+  expectAssignable<DocumentNode>(document)
+  expectAssignable<MercuriusContext>(context)
   return {
     document,
     errors: [
@@ -485,7 +492,9 @@ app.graphql.addHook('preExecution', async function (schema, document, context) {
 })
 
 app.graphql.addHook('preGatewayExecution', async function (schema, document, context) {
-  console.log('preGatewayExecution called')
+  expectAssignable<GraphQLSchema>(schema)
+  expectAssignable<DocumentNode>(document)
+  expectAssignable<MercuriusContext>(context)
   return {
     document,
     errors: [
@@ -494,35 +503,56 @@ app.graphql.addHook('preGatewayExecution', async function (schema, document, con
   }
 })
 
+// GraphQL Request lifecycle hooks
 app.graphql.addHook('onResolution', async function (execution, context) {
-  console.log('onResolution called')
+  expectAssignable<ExecutionResult>(execution)
+  expectAssignable<MercuriusContext>(context)
 })
 
 app.graphql.addHook('preSubscriptionParsing', async function (schema, source, context) {
-  console.log('preSubscriptionParsing called')
+  expectAssignable<GraphQLSchema>(schema)
+  expectAssignable<string>(source)
+  expectAssignable<MercuriusContext>(context)
 })
 
 app.graphql.addHook('preSubscriptionExecution', async function (schema, document, context) {
-  console.log('preSubscriptionExecution called')
+  expectAssignable<GraphQLSchema>(schema)
+  expectAssignable<DocumentNode>(document)
+  expectAssignable<MercuriusContext>(context)
 })
 
 app.graphql.addHook('preGatewaySubscriptionExecution', async function (schema, document, context) {
-  console.log('preGatewaySubscriptionExecution called')
+  expectAssignable<GraphQLSchema>(schema)
+  expectAssignable<DocumentNode>(document)
+  expectAssignable<MercuriusContext>(context)
 })
 
 app.graphql.addHook('onSubscriptionResolution', async function (execution, context) {
-  console.log('onSubscriptionResolution called')
+  expectAssignable<ExecutionResult>(execution)
+  expectAssignable<MercuriusContext>(context)
 })
 
 app.graphql.addHook('onSubscriptionEnd', async function (context) {
-  console.log('onSubscriptionEnd called')
+  expectAssignable<MercuriusContext>(context)
 })
 
 // Hooks containing service metadata
 app.graphql.addHook('preGatewayExecution', async function (schema, document, context, service) {
-  console.log('preGatewayExecution called')
+  expectAssignable<GraphQLSchema>(schema)
+  expectAssignable<DocumentNode>(document)
+  expectAssignable<MercuriusContext>(context)
+  expectAssignable<MercuriusServiceMetadata>(service)
 })
 
 app.graphql.addHook('preGatewaySubscriptionExecution', async function (schema, document, context, service) {
-  console.log('preGatewaySubscriptionExecution called')
+  expectAssignable<GraphQLSchema>(schema)
+  expectAssignable<DocumentNode>(document)
+  expectAssignable<MercuriusContext>(context)
+  expectAssignable<MercuriusServiceMetadata>(service)
+})
+
+// GraphQL Application lifecycle hooks
+app.graphql.addHook('onGatewayReplaceSchema', async function (instance, schema) {
+  expectAssignable<FastifyInstance>(instance)
+  expectAssignable<GraphQLSchema>(schema)
 })