feat(gw): header propagation plugin

.changeset/nervous-kangaroos-smell.md

@@ -0,0 +1,6 @@
+---
+'@graphql-mesh/fusion-runtime': patch
+'@graphql-mesh/serve-runtime': patch
+---
+
+Header Propagation

packages/fusion/runtime/src/utils.ts

@@ -124,6 +124,8 @@ function getTransportExecutor({
   );
 }
 
+export const subgraphNameByExecutionRequest = new WeakMap<ExecutionRequest, string>();
+
 /**
  * This function creates a executor factory that uses the transport packages,
  * and wraps them with the hooks
@@ -145,6 +147,7 @@ export function getOnSubgraphExecute({
 }) {
   const subgraphExecutorMap = new Map<string, Executor>();
   return function onSubgraphExecute(subgraphName: string, executionRequest: ExecutionRequest) {
+    subgraphNameByExecutionRequest.set(executionRequest, subgraphName);
     let executor: Executor = subgraphExecutorMap.get(subgraphName);
     // If the executor is not initialized yet, initialize it
     if (executor == null) {

packages/serve-runtime/src/createGatewayRuntime.ts

@@ -70,6 +70,7 @@ import { useCompleteSubscriptionsOnSchemaChange } from './plugins/useCompleteSub
 import { useContentEncoding } from './plugins/useContentEncoding.js';
 import { useCustomAgent } from './plugins/useCustomAgent.js';
 import { useFetchDebug } from './plugins/useFetchDebug.js';
+import { usePropagateHeaders } from './plugins/usePropagateHeaders.js';
 import { useRequestId } from './plugins/useRequestId.js';
 import { useSubgraphExecuteDebug } from './plugins/useSubgraphExecuteDebug.js';
 import { useUpstreamCancel } from './plugins/useUpstreamCancel.js';
@@ -819,6 +820,10 @@ export function createGatewayRuntime<TContext extends Record<string, any> = Reco
     extraPlugins.push(useHmacUpstreamSignature(config.hmacSignature));
   }
 
+  if (config.propagateHeaders) {
+    extraPlugins.push(usePropagateHeaders(config.propagateHeaders));
+  }
+
   const yoga = createYoga<unknown, GatewayContext & TContext>({
     fetchAPI: config.fetchAPI,
     logging: logger,

packages/serve-runtime/src/index.ts

@@ -3,7 +3,7 @@ export * from './types.js';
 export * from './plugins/useCustomFetch.js';
 export * from './plugins/useStaticFiles.js';
 export * from './getProxyExecutor.js';
-export * from './plugins/useForwardHeaders.js';
+export * from './plugins/usePropagateHeaders.js';
 export * from '@whatwg-node/disposablestack';
 export type { ResolveUserFn, ValidateUserFn } from '@envelop/generic-auth';
 export * from '@graphql-mesh/hmac-upstream-signature';

packages/serve-runtime/src/plugins/usePropagateHeaders.ts

@@ -0,0 +1,81 @@
+import { mapMaybePromise } from '@envelop/core';
+import { subgraphNameByExecutionRequest } from '@graphql-mesh/fusion-runtime';
+import type { TransportEntry } from '@graphql-mesh/transport-common';
+import type { OnFetchHookDone } from '@graphql-mesh/types';
+import type { MaybePromise } from '@graphql-tools/utils';
+import type { GatewayPlugin } from '../types';
+
+interface FromClientToSubgraphsPayload {
+  request: Request;
+  subgraphName: string;
+}
+
+interface FromSubgraphsToClientPayload {
+  response: Response;
+  subgraphName: string;
+}
+
+export interface PropagateHeadersOpts {
+  fromClientToSubgraphs?: (
+    payload: FromClientToSubgraphsPayload,
+  ) => Record<string, string> | void | Promise<Record<string, string | null | undefined> | void>;
+  fromSubgraphsToClient?: (
+    payload: FromSubgraphsToClientPayload,
+  ) => Record<string, string> | void | Promise<Record<string, string | null | undefined> | void>;
+}
+
+export function usePropagateHeaders<TContext>(opts: PropagateHeadersOpts): GatewayPlugin<TContext> {
+  const resHeadersByRequest = new WeakMap<Request, Record<string, string>>();
+  return {
+    onFetch({ executionRequest, context, options, setOptions }) {
+      const subgraphName = subgraphNameByExecutionRequest.get(executionRequest);
+      if (subgraphName != null) {
+        let job: Promise<void> | void;
+        if (opts.fromClientToSubgraphs) {
+          job = mapMaybePromise(
+            opts.fromClientToSubgraphs({
+              request: context.request,
+              subgraphName,
+            }),
+            headers =>
+              setOptions({
+                ...options,
+                headers: {
+                  ...headers,
+                  ...options.headers,
+                },
+              }),
+          );
+        }
+        return mapMaybePromise(job, (): OnFetchHookDone => {
+          if (opts.fromSubgraphsToClient) {
+            return function onFetchDone({ response }) {
+              return mapMaybePromise(
+                opts.fromSubgraphsToClient({
+                  response,
+                  subgraphName,
+                }),
+                headers => {
+                  if (headers) {
+                    resHeadersByRequest.set(context.request, headers);
+                  }
+                },
+              );
+            };
+          }
+        });
+      }
+    },
+    onResponse({ response, request }) {
+      const headers = resHeadersByRequest.get(request);
+      if (headers) {
+        for (const key in headers) {
+          const value = headers[key];
+          if (value) {
+            response.headers.set(key, value);
+          }
+        }
+      }
+    },
+  };
+}

packages/serve-runtime/src/types.ts

@@ -13,6 +13,7 @@ import type { useGenericAuth } from '@envelop/generic-auth';
 import type { Transports, UnifiedGraphPlugin } from '@graphql-mesh/fusion-runtime';
 import type { HMACUpstreamSignatureOptions } from '@graphql-mesh/hmac-upstream-signature';
 import type useMeshResponseCache from '@graphql-mesh/plugin-response-cache';
+import type { PropagateHeadersOpts } from '@graphql-mesh/serve-cli';
 import type { TransportEntry } from '@graphql-mesh/transport-common';
 import type {
   KeyValueCache,
@@ -473,4 +474,9 @@ interface GatewayConfigBase<TContext extends Record<string, any>> {
    * Enable WebHooks handling
    */
   webhooks?: boolean;
+
+  /**
+   * Header Propagation
+   */
+  propagateHeaders?: PropagateHeadersOpts;
 }

packages/serve-runtime/tests/propagateHeaders.spec.ts

@@ -0,0 +1,142 @@
+import { createSchema, createYoga, type Plugin } from 'graphql-yoga';
+import { useCustomFetch } from '@graphql-mesh/serve-runtime';
+import { createGatewayRuntime } from '../src/createGatewayRuntime';
+
+describe('usePropagateHeaders', () => {
+  describe('From Client to the Subgraphs', () => {
+    const requestTrackerPlugin = {
+      onParams: jest.fn((() => {}) as Plugin['onParams']),
+    };
+    const upstream = createYoga({
+      schema: createSchema({
+        typeDefs: /* GraphQL */ `
+          type Query {
+            hello: String
+          }
+        `,
+        resolvers: {
+          Query: {
+            hello: () => 'world',
+          },
+        },
+      }),
+      plugins: [requestTrackerPlugin],
+      logging: !!process.env.DEBUG,
+    });
+    beforeEach(() => {
+      requestTrackerPlugin.onParams.mockClear();
+    });
+    it('forwards specified headers', async () => {
+      await using serveRuntime = createGatewayRuntime({
+        proxy: {
+          endpoint: 'http://localhost:4001/graphql',
+        },
+        propagateHeaders: {
+          fromClientToSubgraphs({ request }) {
+            return {
+              'x-my-header': request.headers.get('x-my-header'),
+              'x-my-other': request.headers.get('x-my-other'),
+            };
+          },
+          fromSubgraphsToClient({ response }) {
+            return {
+              'set-cookies': response.headers.get('set-cookies'),
+            };
+          },
+        },
+        plugins: () => [useCustomFetch(upstream.fetch)],
+        logging: !!process.env.DEBUG,
+      });
+      const response = await serveRuntime.fetch('http://localhost:4000/graphql', {
+        method: 'POST',
+        headers: {
+          'x-my-header': 'my-value',
+          'x-my-other': 'other-value',
+          'x-extra-header': 'extra-value',
+          'content-type': 'application/json',
+        },
+        body: JSON.stringify({
+          query: /* GraphQL */ `
+            query {
+              hello
+            }
+          `,
+          extensions: {
+            randomThing: 'randomValue',
+          },
+        }),
+      });
+
+      const resJson = await response.json();
+      expect(resJson).toEqual({
+        data: {
+          hello: 'world',
+        },
+      });
+
+      // The first call is for the introspection
+      expect(requestTrackerPlugin.onParams).toHaveBeenCalledTimes(2);
+      const onParamsPayload = requestTrackerPlugin.onParams.mock.calls[1][0];
+      // Do not pass extensions
+      expect(onParamsPayload.params.extensions).toBeUndefined();
+      const headersObj = Object.fromEntries(onParamsPayload.request.headers.entries());
+      expect(headersObj['x-my-header']).toBe('my-value');
+      expect(headersObj['x-my-other']).toBe('other-value');
+      expect(headersObj['x-extra-header']).toBeUndefined();
+    });
+    it("forwards specified headers but doesn't override the provided headers", async () => {
+      await using serveRuntime = createGatewayRuntime({
+        logging: !!process.env.DEBUG,
+        proxy: {
+          endpoint: 'http://localhost:4001/graphql',
+          headers: {
+            'x-my-header': 'my-value',
+            'x-extra-header': 'extra-value',
+          },
+        },
+        propagateHeaders: {
+          fromClientToSubgraphs({ request }) {
+            return {
+              'x-my-header': request.headers.get('x-my-header')!,
+              'x-my-other': request.headers.get('x-my-other')!,
+            };
+          },
+        },
+        plugins: () => [useCustomFetch(upstream.fetch)],
+        maskedErrors: false,
+      });
+      const response = await serveRuntime.fetch('http://localhost:4000/graphql', {
+        method: 'POST',
+        headers: {
+          'content-type': 'application/json',
+          'x-my-header': 'my-new-value',
+          'x-my-other': 'other-value',
+        },
+        body: JSON.stringify({
+          query: /* GraphQL */ `
+            query {
+              hello
+            }
+          `,
+        }),
+      });
+
+      const resJson = await response.json();
+      expect(resJson).toEqual({
+        data: {
+          hello: 'world',
+        },
+      });
+
+      // The first call is for the introspection
+      expect(requestTrackerPlugin.onParams).toHaveBeenCalledTimes(2);
+      const onParamsPayload = requestTrackerPlugin.onParams.mock.calls[1][0];
+      // Do not pass extensions
+      expect(onParamsPayload.params.extensions).toBeUndefined();
+      const headersObj = Object.fromEntries(onParamsPayload.request.headers.entries());
+      expect(headersObj['x-my-header']).toBe('my-value');
+      expect(headersObj['x-extra-header']).toBe('extra-value');
+      expect(headersObj['x-my-other']).toBe('other-value');
+    });
+  });
+});