feat: [server] Add new onOperation hook

.changeset/curvy-months-invite.md

@@ -0,0 +1,21 @@
+---
+'graphql-yoga': minor
+---
+
+## New hook `onOperation`
+
+This new hook is called for each GraphQL operaiton to handle. It allows to replace the
+default Yoga operation handler.
+
+Example: Wrap the GraphQL handling pipeline in an `AsyncLocalStorage`
+
+```ts
+function myPlugin(): Plugin {
+  const context = new AsyncLocalStorage();
+  return {
+    onOperation({ operationHandler, setOperationHandler }) {
+      const store = { foo: 'bar' }
+      setOperationHandler((payload) => context.run(store, operationHandler, payload))
+   }
+}
+```

packages/graphql-yoga/__tests__/plugin-hooks.spec.ts

@@ -1,107 +1,153 @@
 import { createSchema, createYoga, type Plugin } from '../src';
 import { eventStream } from './utilities';
 
-test('onParams -> setResult to single execution result', async () => {
-  const plugin: Plugin = {
-    async onParams({ setResult }) {
-      setResult({ data: { hello: 'world' } });
-    },
-  };
-
-  const yoga = createYoga({ plugins: [plugin] });
-
-  const result = await yoga.fetch('http://yoga/graphql', {
-    method: 'POST',
-    body: JSON.stringify({ query: '{ hello }' }),
-    headers: {
-      'Content-Type': 'application/json',
-    },
-  });
-
-  expect(result.status).toBe(200);
-  const body = await result.json();
-  expect(body).toEqual({ data: { hello: 'world' } });
-});
-
-test('onParams -> setResult to event stream execution result', async () => {
-  const plugin: Plugin = {
-    async onParams({ setResult }) {
-      setResult(
-        (async function* () {
-          yield { data: { hello: 'hee' } };
-          yield { data: { hello: 'hoo' } };
-        })(),
-      );
-    },
-  };
-
-  const yoga = createYoga({ plugins: [plugin] });
-
-  const result = await yoga.fetch('http://yoga/graphql', {
-    method: 'POST',
-    body: JSON.stringify({ query: '{ hello }' }),
-    headers: {
-      'Content-Type': 'application/json',
-      Accept: 'text/event-stream',
-    },
-  });
+describe('Plugin hooks', () => {
+  it('context value identity stays the same in all hooks', async () => {
+    const contextValues = [] as Array<unknown>;
+    const yoga = createYoga({
+      schema: createSchema({ typeDefs: `type Query {a:String}` }),
+      plugins: [
+        {
+          onParams(ctx) {
+            contextValues.push(ctx.context);
+          },
+          onParse(ctx) {
+            contextValues.push(ctx.context);
+          },
+          onValidate(ctx) {
+            contextValues.push(ctx.context);
+          },
+          onContextBuilding(ctx) {
+            contextValues.push(ctx.context);
+            // mutate context
+            ctx.extendContext({ a: 1 } as Record<string, unknown>);
+            contextValues.push(ctx.context);
+          },
+          onExecute(ctx) {
+            contextValues.push(ctx.args.contextValue);
+          },
+          onResponse(ctx) {
+            contextValues.push(ctx.serverContext);
+          },
+        } satisfies Plugin,
+      ],
+    });
 
-  expect(result.status).toBe(200);
-  let counter = 0;
-  for await (const value of eventStream(result.body!)) {
-    if (counter === 0) {
-      expect(value).toEqual({ data: { hello: 'hee' } });
-      counter++;
-    } else if (counter === 1) {
-      expect(value).toEqual({ data: { hello: 'hoo' } });
-      counter++;
+    const response = await yoga.fetch('http://localhost/graphql', {
+      method: 'POST',
+      headers: {
+        'content-type': 'application/json',
+      },
+      body: JSON.stringify({ query: '{__typename}' }),
+    });
+    expect(response.status).toEqual(200);
+    expect(await response.json()).toEqual({ data: { __typename: 'Query' } });
+    expect(contextValues).toHaveLength(7);
+    for (const value of contextValues) {
+      expect(value).toBe(contextValues[0]);
     }
-  }
-  expect(counter).toBe(2);
-});
+  });
 
-test('context value identity stays the same in all hooks', async () => {
-  const contextValues = [] as Array<unknown>;
-  const yoga = createYoga({
-    schema: createSchema({ typeDefs: `type Query {a:String}` }),
-    plugins: [
-      {
-        onParams(ctx) {
-          contextValues.push(ctx.context);
-        },
-        onParse(ctx) {
-          contextValues.push(ctx.context);
+  describe('onParams', () => {
+    it('setResult to single execution result', async () => {
+      const plugin: Plugin = {
+        async onParams({ setResult }) {
+          setResult({ data: { hello: 'world' } });
         },
-        onValidate(ctx) {
-          contextValues.push(ctx.context);
-        },
-        onContextBuilding(ctx) {
-          contextValues.push(ctx.context);
-          // mutate context
-          ctx.extendContext({ a: 1 } as Record<string, unknown>);
-          contextValues.push(ctx.context);
+      };
+
+      const yoga = createYoga({ plugins: [plugin] });
+
+      const result = await yoga.fetch('http://yoga/graphql', {
+        method: 'POST',
+        body: JSON.stringify({ query: '{ hello }' }),
+        headers: {
+          'Content-Type': 'application/json',
         },
-        onExecute(ctx) {
-          contextValues.push(ctx.args.contextValue);
+      });
+
+      expect(result.status).toBe(200);
+      const body = await result.json();
+      expect(body).toEqual({ data: { hello: 'world' } });
+    });
+
+    it('setResult to event stream execution result', async () => {
+      const plugin: Plugin = {
+        async onParams({ setResult }) {
+          setResult(
+            (async function* () {
+              yield { data: { hello: 'hee' } };
+              yield { data: { hello: 'hoo' } };
+            })(),
+          );
         },
-        onResponse(ctx) {
-          contextValues.push(ctx.serverContext);
+      };
+
+      const yoga = createYoga({ plugins: [plugin] });
+
+      const result = await yoga.fetch('http://yoga/graphql', {
+        method: 'POST',
+        body: JSON.stringify({ query: '{ hello }' }),
+        headers: {
+          'Content-Type': 'application/json',
+          Accept: 'text/event-stream',
         },
-      } satisfies Plugin,
-    ],
+      });
+
+      expect(result.status).toBe(200);
+      let counter = 0;
+      for await (const value of eventStream(result.body!)) {
+        if (counter === 0) {
+          expect(value).toEqual({ data: { hello: 'hee' } });
+          counter++;
+        } else if (counter === 1) {
+          expect(value).toEqual({ data: { hello: 'hoo' } });
+          counter++;
+        }
+      }
+      expect(counter).toBe(2);
+    });
   });
 
-  const response = await yoga.fetch('http://localhost/graphql', {
-    method: 'POST',
-    headers: {
-      'content-type': 'application/json',
-    },
-    body: JSON.stringify({ query: '{__typename}' }),
+  describe('onOperation', () => {
+    it('should allow to override the graphql operation handling', async () => {
+      const operationHandler = jest.fn().mockReturnValue({ foo: 'bar' });
+      const plugin: Plugin = {
+        onOperation: jest.fn().mockImplementation(({ setOperationHandler }) => {
+          setOperationHandler(operationHandler);
+        }),
+        onEnveloped: jest.fn(),
+      };
+
+      const yoga = createYoga({ plugins: [plugin] });
+
+      const result = await yoga.fetch('http://yoga/graphql', {
+        method: 'POST',
+        body: JSON.stringify({ query: '{ hello }' }),
+        headers: {
+          'Content-Type': 'application/json',
+        },
+      });
+
+      expect(result.status).toBe(200);
+      const body = await result.json();
+      expect(body).toEqual({ foo: 'bar' });
+
+      expect(operationHandler).toHaveBeenCalled();
+      // The entire graphql operation process should have been replaced
+      expect(plugin.onEnveloped).not.toHaveBeenCalled();
+
+      expect(plugin.onOperation).toHaveBeenCalledWith(
+        expect.objectContaining({
+          request: expect.any(Object),
+          params: expect.any(Object),
+          context: expect.objectContaining({
+            waitUntil: expect.any(Function),
+          }),
+          operationHandler: expect.any(Function),
+          setOperationHandler: expect.any(Function),
+        }),
+      );
+    });
   });
-  expect(response.status).toEqual(200);
-  expect(await response.json()).toEqual({ data: { __typename: 'Query' } });
-  expect(contextValues).toHaveLength(7);
-  for (const value of contextValues) {
-    expect(value).toBe(contextValues[0]);
-  }
 });

packages/graphql-yoga/src/plugins/types.ts

@@ -79,6 +79,11 @@ export type Plugin<
      * should be used for sending the result over the wire.
      */
     onResultProcess?: OnResultProcess<TServerContext>;
+    /**
+     * This hook is invoked for each graphql operation.
+     * Here you can set graphql operation handler.
+     */
+    onOperation?: OnOperationHook<TServerContext>;
   };
 
 export type OnYogaInitHook<TServerContext extends Record<string, any>> = (
@@ -203,3 +208,24 @@ export type OnPluginInitEventPayload<PluginContext extends Record<string, any>>
 export type OnPluginInitHook<ContextType extends Record<string, any>> = (
   options: OnPluginInitEventPayload<ContextType>,
 ) => void;
+
+export type OperationHandler<TServerContext extends Record<string, any>> = (payload: {
+  params: GraphQLParams;
+  request: Request;
+  context: TServerContext;
+}) => PromiseOrValue<ExecutionResult | AsyncIterable<ExecutionResult> | undefined>;
+
+export type OnOperationHookPayload<TServerContext extends Record<string, any>> = {
+  context: TServerContext;
+  operationHandler: OperationHandler<TServerContext>;
+  setOperationHandler: (operationHandler: OperationHandler<TServerContext>) => void;
+  params: GraphQLParams;
+  request: Request;
+};
+
+/**
+ * Invoked for each GraphQL operation.
+ */
+export type OnOperationHook<TServerContext extends Record<string, any>> = (
+  payload: OnOperationHookPayload<TServerContext>,
+) => PromiseOrValue<void>;

packages/graphql-yoga/src/server.ts

@@ -52,6 +52,8 @@ import {
   Plugin,
   RequestParser,
   ResultProcessorInput,
+  type OnOperationHook,
+  type OperationHandler,
 } from './plugins/types.js';
 import { GraphiQLOptions, GraphiQLOptionsOrFactory, useGraphiQL } from './plugins/use-graphiql.js';
 import { useHealthCheck } from './plugins/use-health-check.js';
@@ -217,6 +219,7 @@ export class YogaServer<
   >;
   private onRequestParseHooks: OnRequestParseHook<TServerContext>[];
   private onParamsHooks: OnParamsHook<TServerContext>[];
+  private onOperationHooks: OnOperationHook<TServerContext>[];
   private onExecutionResultHooks: OnExecutionResultHook<TServerContext>[];
   private onResultProcessHooks: OnResultProcess<TServerContext>[];
   private maskedErrorsOpts: YogaMaskedErrorOpts | null;
@@ -437,6 +440,7 @@ export class YogaServer<
 
     this.onRequestParseHooks = [];
     this.onParamsHooks = [];
+    this.onOperationHooks = [];
     this.onExecutionResultHooks = [];
     this.onResultProcessHooks = [];
     for (const plugin of this.plugins) {
@@ -452,6 +456,9 @@ export class YogaServer<
         if (plugin.onParams) {
           this.onParamsHooks.push(plugin.onParams);
         }
+        if (plugin.onOperation) {
+          this.onOperationHooks.push(plugin.onOperation);
+        }
         if (plugin.onExecutionResult) {
           this.onExecutionResultHooks.push(plugin.onExecutionResult);
         }
@@ -462,16 +469,7 @@ export class YogaServer<
     }
   }
 
-  async getResultForParams(
-    {
-      params,
-      request,
-    }: {
-      params: GraphQLParams;
-      request: Request;
-    },
-    context: TServerContext,
-  ) {
+  handleOperation: OperationHandler<TServerContext> = async ({ params, request, context }) => {
     let result: ExecutionResult | AsyncIterable<ExecutionResult> | undefined;
 
     try {
@@ -552,6 +550,33 @@ export class YogaServer<
       });
     }
     return result;
+  };
+
+  async getResultForParams(
+    {
+      params,
+      request,
+    }: {
+      params: GraphQLParams;
+      request: Request;
+    },
+    context: TServerContext,
+  ) {
+    let operationHandler = this.handleOperation;
+
+    for (const onOperationHook of this.onOperationHooks) {
+      await onOperationHook({
+        operationHandler,
+        setOperationHandler: newOperationHandler => {
+          operationHandler = newOperationHandler;
+        },
+        context,
+        params,
+        request,
+      });
+    }
+
+    return operationHandler({ params, request, context });
   }
 
   handle: ServerAdapterRequestHandler<TServerContext> = async (