fix: pass request headers for the graph (#1219)

## Summary by Sourcery

Enable forwarding custom request headers through the pagination client
and streamline pagination behavior

New Features:
- Allow passing custom request headers to underlying GraphQL requests
via client.query and pagination

Enhancements:
- Refactor and simplify ListField interface by removing unused fields
- Simplify pagination loop to always continue when a full batch is
returned, eliminating special-case directive tracking

Tests:
- Add tests for passing headers and using request-options object with
@fetchall
- Update existing pagination tests to include headers parameter in mock
calls

Author: janb87

sdk/thegraph/src/utils/pagination.test.ts

@@ -189,7 +189,7 @@ interface RequestVariables {
   };
 }
 
-const requestMock = mock(async (document: DocumentNode, variables: RequestVariables) => {
+const requestMock = mock(async (document: DocumentNode, variables: RequestVariables, _requestHeaders?: HeadersInit) => {
   // Check if the "holders" field in the query has "skip" or "first" arguments set
   let shouldFetchHolders = false;
   let holdersFieldHasPagination = false;
@@ -278,7 +278,6 @@ describe("createTheGraphClientWithPagination", () => {
             symbol
           }
         }`),
-      {},
     );
     expect(result.tokens).toHaveLength(TEST_TOKENS.length);
     expect(result.tokens).toEqual(TEST_TOKENS);
@@ -289,6 +288,26 @@ describe("createTheGraphClientWithPagination", () => {
     );
   });
 
+  it("should return all tokens if @fetchAll is used with a request options object", async () => {
+    const result = await client.query({
+      document: theGraphGraphql(`
+        query {
+          tokens @fetchAll {
+            name
+            symbol
+          }
+        }`),
+      variables: {},
+    });
+    expect(result.tokens).toHaveLength(TEST_TOKENS.length);
+    expect(result.tokens).toEqual(TEST_TOKENS);
+    const expectedCalls = Math.ceil(TEST_TOKENS.length / DEFAULT_PAGE_SIZE);
+    // If the number of tokens is an exact multiple of the page size, we expect one extra call
+    expect(requestMock).toHaveBeenCalledTimes(
+      TEST_TOKENS.length % DEFAULT_PAGE_SIZE === 0 ? expectedCalls + 1 : expectedCalls,
+    );
+  });
+
   it("should use existing pagination variables if they are passed in", async () => {
     const limit = 100;
     const skip = 1;
@@ -476,23 +495,54 @@ describe("createTheGraphClientWithPagination", () => {
         };
       }),
     );
-    expect(requestMock).toHaveBeenNthCalledWith(1, expect.anything(), {
-      orderBy: "name",
-      orderDirection: "asc",
-      where: {
-        name: "Token",
+    expect(requestMock).toHaveBeenNthCalledWith(
+      1,
+      expect.anything(),
+      {
+        orderBy: "name",
+        orderDirection: "asc",
+        where: {
+          name: "Token",
+        },
+        first: DEFAULT_PAGE_SIZE,
+        skip: 0,
       },
-      first: DEFAULT_PAGE_SIZE,
-      skip: 0,
-    });
-    expect(requestMock).toHaveBeenNthCalledWith(2, expect.anything(), {
-      orderBy: "name",
-      orderDirection: "asc",
-      where: {
-        name: "Token",
+      undefined,
+    );
+    expect(requestMock).toHaveBeenNthCalledWith(
+      2,
+      expect.anything(),
+      {
+        orderBy: "name",
+        orderDirection: "asc",
+        where: {
+          name: "Token",
+        },
+        first: DEFAULT_PAGE_SIZE,
+        skip: DEFAULT_PAGE_SIZE,
+      },
+      undefined,
+    );
+  });
+
+  it("should pass request headers to the request", async () => {
+    const result = await client.query(
+      theGraphGraphql(`
+        query {
+          tokens @fetchAll {
+            name
+            symbol
+          }
+        }`),
+      {},
+      {
+        "x-api-key": "test",
       },
-      first: DEFAULT_PAGE_SIZE,
-      skip: DEFAULT_PAGE_SIZE,
+    );
+    expect(result.tokens).toHaveLength(TEST_TOKENS.length);
+    expect(result.tokens).toEqual(TEST_TOKENS);
+    expect(requestMock).toHaveBeenNthCalledWith(1, expect.anything(), expect.anything(), {
+      "x-api-key": "test",
     });
   });
 });

sdk/thegraph/src/utils/pagination.ts

@@ -1,7 +1,7 @@
 import { sortBy } from "es-toolkit";
 import { get, isArray, isEmpty, set } from "es-toolkit/compat";
 import type { TadaDocumentNode } from "gql.tada";
-import { type ArgumentNode, type DocumentNode, Kind, parse, type SelectionNode, visit } from "graphql";
+import { type ArgumentNode, type DocumentNode, Kind, parse, visit } from "graphql";
 import type { GraphQLClient, RequestDocument, RequestOptions, Variables } from "graphql-request";
 
 // Constants for TheGraph limits
@@ -10,16 +10,12 @@ const FIRST_ARG = "first";
 const SKIP_ARG = "skip";
 const FETCH_ALL_DIRECTIVE = "fetchAll";
 
-interface ListField {
+interface ListFieldWithFetchAllDirective {
   path: string[];
   fieldName: string;
-  alias?: string;
   firstValue?: number;
   skipValue?: number;
   otherArgs: ArgumentNode[];
-  selections?: ReadonlyArray<SelectionNode>;
-  hasFetchAllDirective?: boolean;
-  firstValueIsDefault?: boolean; // Track if first value was defaulted
 }
 
 /**
@@ -115,8 +111,8 @@ function extractFetchAllFields(
   document: DocumentNode,
   variables?: Variables,
   fetchAllFields?: Set<string>,
-): ListField[] {
-  const fields: ListField[] = [];
+): ListFieldWithFetchAllDirective[] {
+  const fields: ListFieldWithFetchAllDirective[] = [];
   const pathStack: string[] = [];
 
   visit(document, {
@@ -174,13 +170,9 @@ function extractFetchAllFields(
           fields.push({
             path: [...pathStack],
             fieldName: node.name.value,
-            alias: node.alias?.value,
-            firstValue: hasFetchAllDirective && (firstValue ?? THE_GRAPH_LIMIT),
-            skipValue: hasFetchAllDirective && (skipValue ?? 0),
+            firstValue: firstValue ?? THE_GRAPH_LIMIT,
+            skipValue: skipValue ?? 0,
             otherArgs,
-            selections: node.selectionSet?.selections,
-            hasFetchAllDirective,
-            firstValueIsDefault: hasFetchAllDirective ? firstValue === undefined : false,
           });
         }
       },
@@ -196,7 +188,7 @@ function extractFetchAllFields(
 // Create a query for a single field with specific pagination
 function createSingleFieldQuery(
   document: DocumentNode,
-  targetField: ListField,
+  targetField: ListFieldWithFetchAllDirective,
   skip: number,
   first: number,
 ): DocumentNode {
@@ -251,7 +243,7 @@ function createSingleFieldQuery(
 }
 
 // Create query without list fields
-function createNonListQuery(document: DocumentNode, listFields: ListField[]): DocumentNode | null {
+function createNonListQuery(document: DocumentNode, listFields: ListFieldWithFetchAllDirective[]): DocumentNode | null {
   let hasFields = false;
   const pathStack: string[] = [];
 
@@ -319,7 +311,8 @@ export function createTheGraphClientWithPagination(theGraphClient: Pick<GraphQLC
   async function executeListFieldPagination(
     document: DocumentNode,
     variables: Variables | undefined,
-    field: ListField,
+    field: ListFieldWithFetchAllDirective,
+    requestHeaders?: HeadersInit,
   ): Promise<unknown[]> {
     const results: unknown[] = [];
     let currentSkip = field.skipValue || 0;
@@ -332,11 +325,15 @@ export function createTheGraphClientWithPagination(theGraphClient: Pick<GraphQLC
     while (hasMore) {
       const query = createSingleFieldQuery(document, field, currentSkip, batchSize);
       const existingVariables = filterVariables(variables, query) ?? {};
-      const response = await theGraphClient.request(query, {
-        ...existingVariables,
-        first: batchSize,
-        skip: currentSkip,
-      });
+      const response = await theGraphClient.request(
+        query,
+        {
+          ...existingVariables,
+          first: batchSize,
+          skip: currentSkip,
+        },
+        requestHeaders,
+      );
 
       // Use array path format for es-toolkit's get function
       const data = get(response, field.path) ?? get(response, field.fieldName);
@@ -352,24 +349,8 @@ export function createTheGraphClientWithPagination(theGraphClient: Pick<GraphQLC
       if (isArray(data) && data.length > 0) {
         results.push(...data);
 
-        // Continue fetching if:
-        // 1. We have @fetchAll directive (fetch everything)
-        // 2. We have an explicit first value > THE_GRAPH_LIMIT and haven't reached it
-        // 3. We have a defaulted first value and got a full batch (treating it as "no explicit value")
-        // 4. We have no first value and got a full batch
-        if (field.hasFetchAllDirective) {
-          // With @fetchAll, continue if we got a full batch
-          hasMore = data.length === batchSize;
-        } else if (field.firstValue && !field.firstValueIsDefault) {
-          // With explicit first value (not defaulted), only continue if:
-          // - We haven't reached the requested amount yet
-          // - We got a full batch (indicating more data might exist)
-          hasMore = data.length === batchSize && results.length < field.firstValue;
-        } else {
-          // When first is not specified or was defaulted (using default batch size),
-          // continue if we got a full batch (standard TheGraph pagination behavior)
-          hasMore = data.length === batchSize;
-        }
+        // With @fetchAll, continue if we got a full batch
+        hasMore = data.length === batchSize;
       } else {
         hasMore = false;
       }
@@ -384,16 +365,20 @@ export function createTheGraphClientWithPagination(theGraphClient: Pick<GraphQLC
     async query<TResult, TVariables extends Variables>(
       documentOrOptions: TadaDocumentNode<TResult, TVariables> | RequestDocument | RequestOptions<TVariables, TResult>,
       variablesRaw?: Omit<TVariables, "skip" | "first">,
+      requestHeadersRaw?: HeadersInit,
     ): Promise<TResult> {
       let document: TadaDocumentNode<TResult, TVariables> | RequestDocument;
       let variables: Omit<TVariables, "skip" | "first">;
+      let requestHeaders: HeadersInit | undefined;
 
       if (isRequestOptions(documentOrOptions)) {
         document = documentOrOptions.document;
-        variables = documentOrOptions.variables as TVariables;
+        variables = (documentOrOptions.variables ?? {}) as TVariables;
+        requestHeaders = documentOrOptions.requestHeaders;
       } else {
         document = documentOrOptions;
         variables = variablesRaw ?? ({} as TVariables);
+        requestHeaders = requestHeadersRaw;
       }
 
       // First, detect and strip @fetchAll directives
@@ -404,7 +389,7 @@ export function createTheGraphClientWithPagination(theGraphClient: Pick<GraphQLC
 
       // If no list fields, execute normally
       if (listFields.length === 0) {
-        return theGraphClient.request(processedDocument, variables as Variables);
+        return theGraphClient.request(processedDocument, variables as Variables, requestHeaders);
       }
 
       // Execute paginated queries for all list fields
@@ -416,7 +401,7 @@ export function createTheGraphClientWithPagination(theGraphClient: Pick<GraphQLC
       // Process list fields in parallel for better performance
       const fieldDataPromises = sortedFields.map(async (field) => ({
         field,
-        data: await executeListFieldPagination(processedDocument, variables, field),
+        data: await executeListFieldPagination(processedDocument, variables, field, requestHeaders),
       }));
 
       const fieldResults = await Promise.all(fieldDataPromises);
@@ -434,6 +419,7 @@ export function createTheGraphClientWithPagination(theGraphClient: Pick<GraphQLC
         const nonListResult = await theGraphClient.request(
           nonListQuery,
           filterVariables(variables, nonListQuery) ?? {},
+          requestHeaders,
         );
 
         // Merge results, preserving list data
@@ -447,5 +433,5 @@ export function createTheGraphClientWithPagination(theGraphClient: Pick<GraphQLC
 }
 
 function isRequestOptions(args: unknown): args is RequestOptions<Variables, unknown> {
-  return typeof args === "object" && args !== null && "document" in args && "variables" in args;
+  return typeof args === "object" && args !== null && "document" in args;
 }