From dfd0ca21648fed5cb10ce8e0ceff34d7bd3f6ed3 Mon Sep 17 00:00:00 2001
From: Laurin Quast <laurinquast@googlemail.com>
Date: Thu, 7 Nov 2024 16:15:55 +0100
Subject: [PATCH] refactor: use weak map

---
 .../graphql-yoga/__tests__/batching.spec.ts   | 29 +++++++------------
 packages/graphql-yoga/src/server.ts           | 12 ++++----
 .../src/utils/batch-request-index.ts          |  5 ++++
 3 files changed, 20 insertions(+), 26 deletions(-)
 create mode 100644 packages/graphql-yoga/src/utils/batch-request-index.ts

diff --git a/packages/graphql-yoga/__tests__/batching.spec.ts b/packages/graphql-yoga/__tests__/batching.spec.ts
index ea5741888a..9edb7dc991 100644
--- a/packages/graphql-yoga/__tests__/batching.spec.ts
+++ b/packages/graphql-yoga/__tests__/batching.spec.ts
@@ -1,4 +1,4 @@
-import { createSchema, createYoga, Plugin } from '../src';
+import { createSchema, createYoga, getBatchRequestIndexFromContext, Plugin } from '../src/index.js';
 
 describe('Batching', () => {
   const schema = createSchema({
@@ -296,16 +296,13 @@ describe('Batching', () => {
             expect(params.batchedRequestIndex).toEqual(0);
           },
           onParse(context) {
-            // @ts-expect-error not in types
-            expect(context.context[Symbol.for('yogaBatchedRequestIndex')]).toEqual(0);
+            expect(getBatchRequestIndexFromContext(context.context)).toEqual(0);
           },
           onValidate(context) {
-            // @ts-expect-error not in types
-            expect(context.context[Symbol.for('yogaBatchedRequestIndex')]).toEqual(0);
+            expect(getBatchRequestIndexFromContext(context.context)).toEqual(0);
           },
           onExecute(context) {
-            // @ts-expect-error not in types
-            expect(context.args.contextValue[Symbol.for('yogaBatchedRequestIndex')]).toEqual(0);
+            expect(getBatchRequestIndexFromContext(context.args.contextValue)).toEqual(0);
           },
         } satisfies Plugin,
       ],
@@ -338,31 +335,25 @@ describe('Batching', () => {
           onParse(context) {
             const params = JSON.stringify(context.params);
             if (params === '{"source":"{hello}"}') {
-              // @ts-expect-error not in types
-              expect(context.context[Symbol.for('yogaBatchedRequestIndex')]).toEqual(0);
+              expect(getBatchRequestIndexFromContext(context.context)).toEqual(0);
             } else if (params === '{"source":"{bye}"}') {
-              // @ts-expect-error not in types
-              expect(context.context[Symbol.for('yogaBatchedRequestIndex')]).toEqual(1);
+              expect(getBatchRequestIndexFromContext(context.context)).toEqual(1);
             }
           },
           onValidate(context) {
             const params = JSON.stringify(context.params);
             if (params === '{"source":"{hello}"}') {
-              // @ts-expect-error not in types
-              expect(context.context[Symbol.for('yogaBatchedRequestIndex')]).toEqual(0);
+              expect(getBatchRequestIndexFromContext(context.context)).toEqual(0);
             } else if (params === '{"source":"{bye}"}') {
-              // @ts-expect-error not in types
-              expect(context.context[Symbol.for('yogaBatchedRequestIndex')]).toEqual(1);
+              expect(getBatchRequestIndexFromContext(context.context)).toEqual(1);
             }
           },
           onExecute(context) {
             const params = JSON.stringify(context.args.contextValue.params);
             if (params === '{"source":"{hello}"}') {
-              // @ts-expect-error not in types
-              expect(context.context[Symbol.for('yogaBatchedRequestIndex')]).toEqual(0);
+              expect(getBatchRequestIndexFromContext(context.args.contextValue)).toEqual(0);
             } else if (params === '{"source":"{bye}"}') {
-              // @ts-expect-error not in types
-              expect(context.context[Symbol.for('yogaBatchedRequestIndex')]).toEqual(1);
+              expect(getBatchRequestIndexFromContext(context.args.contextValue)).toEqual(1);
             }
           },
         } satisfies Plugin,
diff --git a/packages/graphql-yoga/src/server.ts b/packages/graphql-yoga/src/server.ts
index 29f4f11419..3bb1088730 100644
--- a/packages/graphql-yoga/src/server.ts
+++ b/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';
 
 /**
@@ -498,18 +499,15 @@ export class YogaServer<
                 params,
               };
 
-        let batchIndexPartial: object = {};
-
-        if (batchedRequestIndex !== undefined) {
-          batchIndexPartial = { [Symbol.for('yogaBatchedRequestIndex')]: batchedRequestIndex };
-        }
-
         context = Object.assign(
           batchedRequestIndex === undefined ? serverContext : Object.create(serverContext),
           additionalContext,
-          batchIndexPartial,
         );
 
+        if (batchedRequestIndex !== undefined) {
+          batchRequestIndexMap.set(context, batchedRequestIndex);
+        }
+
         const enveloped = this.getEnveloped(context);
 
         this.logger.debug(`Processing GraphQL Parameters`);
diff --git a/packages/graphql-yoga/src/utils/batch-request-index.ts b/packages/graphql-yoga/src/utils/batch-request-index.ts
new file mode 100644
index 0000000000..ba5bca787d
--- /dev/null
+++ b/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;
+}