feat: batching index in phases

.changeset/young-lamps-juggle.md

@@ -0,0 +1,36 @@
+---
+'graphql-yoga': minor
+---
+
+Support accessing the index of a batched request within the plugin phases.
+
+The helper function `getBatchRequestIndexFromContext` can be used for getting the current batch
+requests index for the ongoing execution.
+
+```ts
+import { createYoga, getBatchRequestIndexFromContext, Plugin } from 'graphql-yoga'
+
+const yoga = createYoga({
+  batched: true,
+  plugins: [
+    {
+      onParams(params) {
+        // undefined or number
+        console.log(params.batchedRequestIndex)
+      },
+      onParse(context) {
+        // undefined or number
+        console.log(getBatchRequestIndexFromContext(context.context))
+      },
+      onValidate(context) {
+        // undefined or number
+        console.log(getBatchRequestIndexFromContext(context.context))
+      },
+      onExecute(context) {
+        // undefined or number
+        console.log(getBatchRequestIndexFromContext(context.args.contextValue))
+      }
+    } satisfies Plugin
+  ]
+})
+```

packages/graphql-yoga/__tests__/batching.spec.ts

@@ -1,5 +1,4 @@
-import { createSchema } from '../src/schema';
-import { createYoga } from '../src/server';
+import { createSchema, createYoga, getBatchRequestIndexFromContext, Plugin } from '../src/index.js';
 
 describe('Batching', () => {
   const schema = createSchema({
@@ -287,4 +286,88 @@ describe('Batching', () => {
       ],
     });
   });
+  it('batching index is forwarded into hooks (single batch)', async () => {
+    const yoga = createYoga({
+      schema,
+      batching: true,
+      plugins: [
+        {
+          onParams(params) {
+            expect(params.batchedRequestIndex).toEqual(0);
+          },
+          onParse(context) {
+            expect(getBatchRequestIndexFromContext(context.context)).toEqual(0);
+          },
+          onValidate(context) {
+            expect(getBatchRequestIndexFromContext(context.context)).toEqual(0);
+          },
+          onExecute(context) {
+            expect(getBatchRequestIndexFromContext(context.args.contextValue)).toEqual(0);
+          },
+        } satisfies Plugin,
+      ],
+    });
+
+    const response = await yoga.fetch('http://yoga.com/graphql', {
+      method: 'POST',
+      headers: {
+        accept: 'application/graphql-response+json',
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify([{ query: '{hello}' }]),
+    });
+    expect(await response.text()).toEqual(`[{"data":{"hello":"hello"}}]`);
+  });
+  it('batching index is forwarded into hooks (multiple batches)', async () => {
+    const yoga = createYoga({
+      schema,
+      batching: true,
+      plugins: [
+        {
+          onParams(context) {
+            const params = JSON.stringify(context.params);
+            if (params === '{"source":"{hello}"}') {
+              expect(context.batchedRequestIndex).toEqual(0);
+            } else if (params === '{"source":"{bye}"}') {
+              expect(context.batchedRequestIndex).toEqual(1);
+            }
+          },
+          onParse(context) {
+            const params = JSON.stringify(context.params);
+            if (params === '{"source":"{hello}"}') {
+              expect(getBatchRequestIndexFromContext(context.context)).toEqual(0);
+            } else if (params === '{"source":"{bye}"}') {
+              expect(getBatchRequestIndexFromContext(context.context)).toEqual(1);
+            }
+          },
+          onValidate(context) {
+            const params = JSON.stringify(context.params);
+            if (params === '{"source":"{hello}"}') {
+              expect(getBatchRequestIndexFromContext(context.context)).toEqual(0);
+            } else if (params === '{"source":"{bye}"}') {
+              expect(getBatchRequestIndexFromContext(context.context)).toEqual(1);
+            }
+          },
+          onExecute(context) {
+            const params = JSON.stringify(context.args.contextValue.params);
+            if (params === '{"source":"{hello}"}') {
+              expect(getBatchRequestIndexFromContext(context.args.contextValue)).toEqual(0);
+            } else if (params === '{"source":"{bye}"}') {
+              expect(getBatchRequestIndexFromContext(context.args.contextValue)).toEqual(1);
+            }
+          },
+        } satisfies Plugin,
+      ],
+    });
+
+    const response = await yoga.fetch('http://yoga.com/graphql', {
+      method: 'POST',
+      headers: {
+        accept: 'application/graphql-response+json',
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify([{ query: '{hello}' }, { query: '{bye}' }]),
+    });
+    expect(await response.text()).toEqual(`[{"data":{"hello":"hello"}},{"data":{"bye":"bye"}}]`);
+  });
 });

packages/graphql-yoga/src/index.ts

@@ -12,6 +12,7 @@ export * from './types.js';
 export { maskError } from './utils/mask-error.js';
 export { type OnParamsEventPayload } from './plugins/types.js';
 export { createLRUCache } from './utils/create-lru-cache.js';
+export { getBatchRequestIndexFromContext } from './utils/batch-request-index.js';
 export { mergeSchemas } from '@graphql-tools/schema';
 export {
   // Handy type utils

packages/graphql-yoga/src/plugins/types.ts

@@ -113,6 +113,8 @@ export interface OnParamsEventPayload {
   setParams: (params: GraphQLParams) => void;
   setResult: (result: ExecutionResult | AsyncIterable<ExecutionResult>) => void;
   fetchAPI: FetchAPI;
+  /** Index of the batched request if it is a batched request */
+  batchedRequestIndex?: number;
 }
 
 export type OnResultProcess<TServerContext> = (

packages/graphql-yoga/src/server.ts

@@ -70,6 +70,7 @@ import {
   YogaInitialContext,
   YogaMaskedErrorOpts,
 } from './types.js';
+import { batchRequestIndexMap } from './utils/batch-request-index.js';
 import { maskError } from './utils/mask-error.js';
 
 /**
@@ -462,13 +463,12 @@ export class YogaServer<
     {
       params,
       request,
-      batched,
     }: {
       params: GraphQLParams;
       request: Request;
-      batched: boolean;
     },
     serverContext: TServerContext,
+    batchedRequestIndex?: number,
   ) {
     let result: ExecutionResult | AsyncIterable<ExecutionResult> | undefined;
     let context = serverContext as TServerContext & YogaInitialContext;
@@ -484,6 +484,7 @@ export class YogaServer<
             result = newResult;
           },
           fetchAPI: this.fetchAPI,
+          batchedRequestIndex,
         });
       }
 
@@ -499,10 +500,14 @@ export class YogaServer<
               };
 
         context = Object.assign(
-          batched ? Object.create(serverContext) : serverContext,
+          batchedRequestIndex === undefined ? serverContext : Object.create(serverContext),
           additionalContext,
         );
 
+        if (batchedRequestIndex !== undefined) {
+          batchRequestIndexMap.set(context, batchedRequestIndex);
+        }
+
         const enveloped = this.getEnveloped(context);
 
         this.logger.debug(`Processing GraphQL Parameters`);
@@ -605,22 +610,21 @@ export class YogaServer<
 
     const result = (await (Array.isArray(requestParserResult)
       ? Promise.all(
-          requestParserResult.map(params =>
+          requestParserResult.map((params, index) =>
             this.getResultForParams(
               {
                 params,
                 request,
-                batched: true,
               },
               serverContext,
+              index,
             ),
           ),
         )
       : this.getResultForParams(
           {
             params: requestParserResult,
             request,
-            batched: false,
           },
           serverContext,
         ))) as ResultProcessorInput;

packages/graphql-yoga/src/utils/batch-request-index.ts

@@ -0,0 +1,5 @@
+export const batchRequestIndexMap = new WeakMap<object, number>();
+
+export function getBatchRequestIndexFromContext(context: object): number | null {
+  return batchRequestIndexMap.get(context) ?? null;
+}