From 9c889ed01f42be43e24f0134285c8cd5a2f2179d Mon Sep 17 00:00:00 2001
From: Valentin Cocaud <v.cocaud@gmail.com>
Date: Wed, 5 Feb 2025 15:23:28 +0100
Subject: [PATCH] feat: [server] Add new  hook

---
 packages/graphql-yoga/src/plugins/types.ts | 26 +++++++++++++
 packages/graphql-yoga/src/server.ts        | 45 +++++++++++++++++-----
 2 files changed, 61 insertions(+), 10 deletions(-)

diff --git a/packages/graphql-yoga/src/plugins/types.ts b/packages/graphql-yoga/src/plugins/types.ts
index e29dd75399..412c08f35c 100644
--- a/packages/graphql-yoga/src/plugins/types.ts
+++ b/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>;
diff --git a/packages/graphql-yoga/src/server.ts b/packages/graphql-yoga/src/server.ts
index f216d849f4..5e4150aeef 100644
--- a/packages/graphql-yoga/src/server.ts
+++ b/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 (