feat: uniquely identify batched queries (#849)

mercurius-js · Aug 29, 2022 · 911a5ef · 911a5ef
1 parent 27e9df4
commit 911a5ef
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 47 deletions.
diff --git a/index.d.ts b/index.d.ts
@@ -25,6 +25,9 @@ export interface PubSub {
 export interface MercuriusContext {
   app: FastifyInstance;
   reply: FastifyReply;
+  operationsCount?: number;
+  operationId?: number;
+  __currentQuery: string;
   /**
    * __Caution__: Only available if `subscriptions` are enabled
    */

diff --git a/lib/routes.js b/lib/routes.js
@@ -352,13 +352,23 @@ module.exports = async function (app, opts) {
 
     if (allowBatchedQueries && Array.isArray(request.body)) {
       // Batched query
-      return Promise.all(request.body.map(r =>
-        execute(r, request, reply)
-          .catch(e => {
-            const { response } = errorFormatter(e, request[kRequestContext])
-            return response
-          })
-      ))
+      const operationsCount = request.body.length
+
+      Object.assign(request[kRequestContext], { operationsCount })
+
+      return Promise.all(request.body.map(async (r, operationId) => {
+        // Create individual reqs for multiple operations, otherwise reference the original req
+        const operationReq = operationsCount > 1 ? Object.create(request) : request
+
+        Object.assign(operationReq[kRequestContext], { operationId })
+
+        try {
+          return await execute(r, operationReq, reply)
+        } catch (e) {
+          const { response } = errorFormatter(e, request[kRequestContext])
+          return response
+        }
+      }))
     } else {
       // Regular query
       return execute(request.body, request, reply)

diff --git a/test/batched.js b/test/batched.js
@@ -2,6 +2,7 @@
 
 const { test } = require('tap')
 const Fastify = require('fastify')
+const sinon = require('sinon')
 const GQL = require('..')
 
 test('POST regular query', async (t) => {
@@ -346,3 +347,44 @@ test('POST batched query with a resolver which succeeds and a resolver which thr
 
   t.same(JSON.parse(res.body), [{ data: { add: 3 } }, { data: null, errors: [{ message: 'Internal Server Error' }] }])
 })
+
+test('POST batched query has an individual context for each operation', async (t) => {
+  const app = Fastify()
+
+  const contextSpy = sinon.spy()
+
+  const schema = `
+      type Query {
+        test: String
+      }
+    `
+
+  const resolvers = {
+    test: (_, ctx) => contextSpy(ctx.operationId, ctx.operationsCount, ctx.__currentQuery)
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    allowBatchedQueries: true
+  })
+
+  await app.inject({
+    method: 'POST',
+    url: '/graphql',
+    body: [
+      {
+        operationName: 'TestQuery',
+        query: 'query TestQuery { test }'
+      },
+      {
+        operationName: 'DoubleQuery',
+        query: 'query DoubleQuery { test }'
+      }
+    ]
+  })
+
+  sinon.assert.calledTwice(contextSpy)
+  sinon.assert.calledWith(contextSpy, 0, 2, sinon.match(/TestQuery/))
+  sinon.assert.calledWith(contextSpy, 1, 2, sinon.match(/DoubleQuery/))
+})
diff --git a/test/options.js b/test/options.js
@@ -30,8 +30,6 @@ const resolvers = {
 }
 
 test('call compileQuery with correct options if compilerOptions specified', async t => {
-  t.plan(1)
-
   const app = Fastify()
   t.teardown(() => app.close())
 
@@ -82,5 +80,5 @@ test('call compileQuery with correct options if compilerOptions specified', asyn
     body: JSON.stringify({ query })
   })
 
-  t.ok(compileQueryStub.calledOnceWith(sinon.match.any, sinon.match.any, sinon.match.any, { customJSONSerializer: true }))
+  sinon.assert.calledOnceWithExactly(compileQueryStub, sinon.match.any, sinon.match.any, sinon.match.any, { customJSONSerializer: true })
 })
diff --git a/test/subscription-connection.js b/test/subscription-connection.js
@@ -637,24 +637,17 @@ test('subscription connection handles query when fullWsTransport: true', async (
     })
   )
 
-  t.ok(
-    send.withArgs(
-      JSON.stringify({
-        type: 'next',
-        id: 1,
-        payload: {}
-      })
-    ).calledOnce
-  )
-  t.ok(
-    send.withArgs(
-      JSON.stringify({
-        type: 'complete',
-        id: 1,
-        payload: null
-      })
-    ).calledOnce
-  )
+  sinon.assert.calledTwice(send)
+  sinon.assert.calledWith(send, JSON.stringify({
+    type: 'next',
+    id: 1,
+    payload: {}
+  }))
+  sinon.assert.calledWith(send, JSON.stringify({
+    type: 'complete',
+    id: 1,
+    payload: null
+  }))
 })
 
 test('subscription connection handles mutation when fullWsTransport: true', async (t) => {
@@ -688,24 +681,17 @@ test('subscription connection handles mutation when fullWsTransport: true', asyn
     })
   )
 
-  t.ok(
-    send.withArgs(
-      JSON.stringify({
-        type: 'next',
-        id: 1,
-        payload: {}
-      })
-    ).calledOnce
-  )
-  t.ok(
-    send.withArgs(
-      JSON.stringify({
-        type: 'complete',
-        id: 1,
-        payload: null
-      })
-    ).calledOnce
-  )
+  sinon.assert.calledTwice(send)
+  sinon.assert.calledWith(send, JSON.stringify({
+    type: 'next',
+    id: 1,
+    payload: {}
+  }))
+  sinon.assert.calledWith(send, JSON.stringify({
+    type: 'complete',
+    id: 1,
+    payload: null
+  }))
 })
 
 test('subscription data is released right after it ends', async (t) => {

diff --git a/test/types/index.ts b/test/types/index.ts
@@ -1,4 +1,4 @@
-import { expectAssignable, expectError } from 'tsd'
+import { expectAssignable, expectError, expectType } from 'tsd'
 /* eslint-disable no-unused-expressions */
 import { EventEmitter } from 'events'
 // eslint-disable-next-line no-unused-vars
@@ -780,3 +780,10 @@ expectError(() => {
 expectError(() => {
   return mercurius.defaultErrorFormatter({}, undefined)
 })
+
+// Context contains correct information about (batched) query identity
+app.graphql.addHook('onResolution', async function (_execution, context) {
+  expectType<number | undefined>(context.operationId)
+  expectType<number | undefined>(context.operationsCount)
+  expectType<string>(context.__currentQuery)
+})