fix: batch delegation collapses two requests with different arguments into a single memoized loader (#382)

Author: ardatan

.changeset/nervous-cups-scream.md

@@ -0,0 +1,7 @@
+---
+'@graphql-tools/batch-delegate': patch
+---
+
+Memoize the batched delegation loader by respecting the arguments correctly
+
+[See more details in the PR](https://github.com/ardatan/graphql-tools/pull/5189)

packages/batch-delegate/src/getLoader.ts

@@ -117,6 +117,11 @@ export function getLoader<K = any, V = any, C = K>(
     cacheKey += memoizedPrint(selectionSet);
   }
 
+  const fieldNode = fieldNodes?.[0];
+  if (fieldNode?.arguments) {
+    cacheKey += fieldNode.arguments.map((arg) => memoizedPrint(arg)).join(',');
+  }
+
   let loader = loaders.get(cacheKey);
 
   if (loader === undefined) {

packages/batch-delegate/tests/basic.example.test.ts

@@ -366,4 +366,92 @@ describe('batch delegation within basic stitching example', () => {
       ],
     });
   });
+  test('uses a single call even when delegating the same field multiple times', async () => {
+    let numCalls = 0;
+
+    const chirpSchema = makeExecutableSchema({
+      typeDefs: /* GraphQL */ `
+        type Chirp {
+          chirpedAtUserId: ID!
+        }
+        type Query {
+          trendingChirps: [Chirp]
+        }
+      `,
+      resolvers: {
+        Query: {
+          trendingChirps: () => [
+            { chirpedAtUserId: 1 },
+            { chirpedAtUserId: 2 },
+          ],
+        },
+      },
+    });
+
+    // Mocked author schema
+    const authorSchema = makeExecutableSchema({
+      typeDefs: /* GraphQL */ `
+        type User {
+          email: String
+        }
+        type Query {
+          usersByIds(ids: [ID!]): [User]
+        }
+      `,
+      resolvers: {
+        Query: {
+          usersByIds: (_root, args) => {
+            numCalls++;
+            return args.ids.map((id: string) => ({ email: `${id}@test.com` }));
+          },
+        },
+      },
+    });
+
+    const linkTypeDefs = /* GraphQL */ `
+      extend type Chirp {
+        chirpedAtUser: User
+      }
+    `;
+
+    const stitchedSchema = stitchSchemas({
+      subschemas: [chirpSchema, authorSchema],
+      typeDefs: linkTypeDefs,
+      resolvers: {
+        Chirp: {
+          chirpedAtUser: {
+            selectionSet: `{ chirpedAtUserId }`,
+            resolve(chirp, _args, context, info) {
+              return batchDelegateToSchema({
+                schema: authorSchema,
+                operation: 'query' as OperationTypeNode,
+                fieldName: 'usersByIds',
+                key: chirp.chirpedAtUserId,
+                argsFromKeys: (ids) => ({ ids }),
+                context,
+                info,
+              });
+            },
+          },
+        },
+      },
+    });
+
+    const query = /* GraphQL */ `
+      query {
+        trendingChirps {
+          chirp1: chirpedAtUser {
+            email
+          }
+          chirp2: chirpedAtUser {
+            email
+          }
+        }
+      }
+    `;
+
+    await execute({ schema: stitchedSchema, document: parse(query) });
+
+    expect(numCalls).toEqual(1);
+  });
 });

packages/batch-delegate/tests/cacheByArguments.test.ts

@@ -0,0 +1,113 @@
+import { batchDelegateToSchema } from '@graphql-tools/batch-delegate';
+import { execute, isIncrementalResult } from '@graphql-tools/executor';
+import { makeExecutableSchema } from '@graphql-tools/schema';
+import { stitchSchemas } from '@graphql-tools/stitch';
+import { OperationTypeNode, parse } from 'graphql';
+import { describe, expect, test } from 'vitest';
+
+describe('non-key arguments are taken into account when memoizing result', () => {
+  test('memoizes non-key arguments as part of batch delegation', async () => {
+    let numCalls = 0;
+
+    const chirpSchema = makeExecutableSchema({
+      typeDefs: /* GraphQL */ `
+        type Chirp {
+          chirpedAtUserId: ID!
+        }
+        type Query {
+          trendingChirps: [Chirp]
+        }
+      `,
+      resolvers: {
+        Query: {
+          trendingChirps: () => [
+            { chirpedAtUserId: 1 },
+            { chirpedAtUserId: 2 },
+          ],
+        },
+      },
+    });
+
+    // Mocked author schema
+    const authorSchema = makeExecutableSchema({
+      typeDefs: /* GraphQL */ `
+        type User {
+          email: String!
+        }
+        type Query {
+          usersByIds(ids: [ID!], obfuscateEmail: Boolean!): [User]
+        }
+      `,
+      resolvers: {
+        Query: {
+          usersByIds: (_root, args) => {
+            numCalls++;
+            return args.ids.map((id: string) => ({
+              email: args.obfuscateEmail ? '***' : `${id}@test.com`,
+            }));
+          },
+        },
+      },
+    });
+
+    const linkTypeDefs = /* GraphQL */ `
+      extend type Chirp {
+        chirpedAtUser(obfuscateEmail: Boolean!): User
+      }
+    `;
+
+    const stitchedSchema = stitchSchemas({
+      subschemas: [chirpSchema, authorSchema],
+      typeDefs: linkTypeDefs,
+      resolvers: {
+        Chirp: {
+          chirpedAtUser: {
+            selectionSet: `{ chirpedAtUserId }`,
+            resolve(chirp, args, context, info) {
+              return batchDelegateToSchema({
+                schema: authorSchema,
+                operation: 'query' as OperationTypeNode,
+                fieldName: 'usersByIds',
+                key: chirp.chirpedAtUserId,
+                argsFromKeys: (ids) => ({ ids, ...args }),
+                context,
+                info,
+              });
+            },
+          },
+        },
+      },
+    });
+
+    const query = /* GraphQL */ `
+      query {
+        trendingChirps {
+          withObfuscatedEmail: chirpedAtUser(obfuscateEmail: true) {
+            email
+          }
+          withoutObfuscatedEmail: chirpedAtUser(obfuscateEmail: false) {
+            email
+          }
+        }
+      }
+    `;
+
+    const result = await execute({
+      schema: stitchedSchema,
+      document: parse(query),
+    });
+
+    expect(numCalls).toEqual(2);
+
+    if (isIncrementalResult(result)) throw Error('result is incremental');
+
+    expect(result.errors).toBeUndefined();
+
+    const chirps: any = result.data!['trendingChirps'];
+    expect(chirps[0].withObfuscatedEmail.email).toBe(`***`);
+    expect(chirps[1].withObfuscatedEmail.email).toBe(`***`);
+
+    expect(chirps[0].withoutObfuscatedEmail.email).toBe(`[email protected]`);
+    expect(chirps[1].withoutObfuscatedEmail.email).toBe(`[email protected]`);
+  });
+});