refactor: use weak map

dotansimha · Nov 7, 2024 · dfd0ca2 · dfd0ca2
1 parent ebc06c9
commit dfd0ca2
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 26 deletions.
diff --git 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
@@ -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
@@ -0,0 +1,5 @@
+export const batchRequestIndexMap = new WeakMap<object, number>();
+
+export function getBatchRequestIndexFromContext(context: object): number | null {
+  return batchRequestIndexMap.get(context) ?? null;
+}