feat: new onQueryPlan hook

.changeset/four-ads-yell.md

@@ -0,0 +1,7 @@
+---
+'@graphql-mesh/fusion-runtime': minor
+'@graphql-hive/router-runtime': minor
+'@graphql-hive/gateway-runtime': minor
+---
+
+New `onQueryPlan` hook to handle query planning with Rust QP

packages/fusion-runtime/src/unifiedGraphManager.ts

@@ -31,6 +31,7 @@ import {
 import { usingHiveRouterRuntime } from '~internal/env';
 import type { DocumentNode, GraphQLError, GraphQLSchema } from 'graphql';
 import { buildASTSchema, buildSchema, isSchema } from 'graphql';
+import { OnQueryPlanHook } from '../../router-runtime/src/types';
 import { handleFederationSupergraph as stitchingUnifiedGraphHandler } from './federation/supergraph';
 import {
   compareSchemas,
@@ -84,6 +85,7 @@ export interface UnifiedGraphHandlerOpts {
   onDelegationPlanHooks?: OnDelegationPlanHook<any>[];
   onDelegationStageExecuteHooks?: OnDelegationStageExecuteHook<any>[];
   onDelegateHooks?: OnDelegateHook<unknown>[];
+  onQueryPlanHooks?: OnQueryPlanHook<any>[];
   /**
    * Configure the batch delegation options for all merged types in all subschemas.
    */
@@ -119,6 +121,7 @@ export interface UnifiedGraphManagerOptions<TContext> {
   onDelegateHooks?: OnDelegateHook<unknown>[];
   onDelegationPlanHooks?: OnDelegationPlanHook<TContext>[];
   onDelegationStageExecuteHooks?: OnDelegationStageExecuteHook<TContext>[];
+  onQueryPlanHooks?: OnQueryPlanHook<TContext>[];
   /**
    * Whether to batch the subgraph executions.
    * @default true
@@ -167,6 +170,7 @@ export class UnifiedGraphManager<TContext> implements AsyncDisposable {
   private onSubgraphExecuteHooks: OnSubgraphExecuteHook<TContext>[];
   private onDelegationPlanHooks: OnDelegationPlanHook<TContext>[];
   private onDelegationStageExecuteHooks: OnDelegationStageExecuteHook<TContext>[];
+  private onQueryPlanHooks: OnQueryPlanHook<TContext>[];
   private inContextSDK: any;
   private initialUnifiedGraph$?: MaybePromise<GraphQLSchema>;
   private polling$?: MaybePromise<void>;
@@ -189,6 +193,7 @@ export class UnifiedGraphManager<TContext> implements AsyncDisposable {
     this.onDelegationPlanHooks = opts?.onDelegationPlanHooks || [];
     this.onDelegationStageExecuteHooks =
       opts?.onDelegationStageExecuteHooks || [];
+    this.onQueryPlanHooks = opts?.onQueryPlanHooks || [];
     if (opts.pollingInterval != null) {
       opts.transportContext?.log.debug(
         `Starting polling to Supergraph with interval ${millisecondsToStr(opts.pollingInterval)}`,
@@ -355,6 +360,7 @@ export class UnifiedGraphManager<TContext> implements AsyncDisposable {
         onDelegationPlanHooks: this.onDelegationPlanHooks,
         onDelegationStageExecuteHooks: this.onDelegationStageExecuteHooks,
         onDelegateHooks: this.opts.onDelegateHooks,
+        onQueryPlanHooks: this.onQueryPlanHooks,
         batchDelegateOptions: this.opts.batchDelegateOptions,
         log: this.opts.transportContext?.log,
         handleProgressiveOverride: this.opts.handleProgressiveOverride

packages/fusion-runtime/src/utils.ts

@@ -42,6 +42,7 @@ import {
   type GraphQLSchema,
 } from 'graphql';
 import type { GraphQLOutputType, GraphQLResolveInfo } from 'graphql/type';
+import { OnQueryPlanHook } from '../../router-runtime/src/types';
 import { restoreExtraDirectives } from './federation/supergraph';
 import {
   Instrumentation,
@@ -428,6 +429,7 @@ export interface UnifiedGraphPlugin<TContext> {
   onSubgraphExecute?: OnSubgraphExecuteHook<TContext>;
   onDelegationPlan?: OnDelegationPlanHook<TContext>;
   onDelegationStageExecute?: OnDelegationStageExecuteHook<TContext>;
+  onQueryPlan?: OnQueryPlanHook<TContext>;
 }
 
 export type OnSubgraphExecuteHook<TContext = any> = (

packages/router-runtime/src/handler.ts

@@ -7,19 +7,20 @@ import {
 import { createDefaultExecutor } from '@graphql-mesh/transport-common';
 import { defaultPrintFn } from '@graphql-tools/executor-common';
 import { filterInternalFieldsAndTypes } from '@graphql-tools/federation';
-import { ExecutionResult, isAsyncIterable } from '@graphql-tools/utils';
+import {
+  ExecutionRequest,
+  ExecutionResult,
+  isAsyncIterable,
+} from '@graphql-tools/utils';
 import {
   handleMaybePromise,
   mapAsyncIterator,
   MaybePromise,
 } from '@whatwg-node/promise-helpers';
 import { BREAK, DocumentNode, visit } from 'graphql';
 import { executeQueryPlan } from './executor';
-import {
-  getLazyFactory,
-  getLazyValue,
-  queryPlanForExecutionRequestContext,
-} from './utils';
+import { getLazyFactory, getLazyValue } from './utils';
+import { wrapQueryPlanFnWithHooks } from './wrapQueryPlanFnWithHooks';
 
 export function unifiedGraphHandler(
   opts: UnifiedGraphHandlerOpts,
@@ -47,13 +48,17 @@ export function unifiedGraphHandler(
     DocumentNode,
     Map<string | null, MaybePromise<QueryPlan>>
   >();
-  function planDocument(document: DocumentNode, operationName: string | null) {
+  function defaultQueryPlanFn({
+    document,
+    operationName,
+  }: ExecutionRequest): MaybePromise<QueryPlan> {
     let operationCache = documentOperationPlanCache.get(document);
 
     // we dont need to worry about releasing values. the map values in weakmap
     // will all be released when document node is GCed
+    const operationNameKey = operationName || null;
     if (operationCache) {
-      const plan = operationCache.get(operationName);
+      const plan = operationCache.get(operationNameKey);
       if (plan) {
         return plan;
       }
@@ -64,14 +69,18 @@ export function unifiedGraphHandler(
 
     const plan = handleMaybePromise(getQueryPlanner, (qp) =>
       qp.plan(defaultPrintFn(document), operationName).then((queryPlan) => {
-        operationCache.set(operationName, queryPlan);
+        operationCache.set(operationNameKey, queryPlan);
         return queryPlan;
       }),
     );
-    operationCache.set(operationName, plan);
+    operationCache.set(operationNameKey, plan);
     return plan;
   }
 
+  const wrappedQueryPlanFn = opts.onQueryPlanHooks?.length
+    ? wrapQueryPlanFnWithHooks(defaultQueryPlanFn, opts.onQueryPlanHooks)
+    : defaultQueryPlanFn;
+
   return {
     unifiedGraph: supergraphSchema,
     getSubgraphSchema(subgraphName: string) {
@@ -82,18 +91,9 @@ export function unifiedGraphHandler(
         return defaultExecutor(executionRequest);
       }
       return handleMaybePromise(
-        () =>
-          planDocument(
-            executionRequest.document,
-            executionRequest.operationName || null,
-          ),
-        (queryPlan) => {
-          queryPlanForExecutionRequestContext.set(
-            // setter like getter
-            executionRequest.context || executionRequest.document,
-            queryPlan,
-          );
-          return executeQueryPlan({
+        () => wrappedQueryPlanFn(executionRequest),
+        (queryPlan) =>
+          executeQueryPlan({
             supergraphSchema,
             executionRequest,
             onSubgraphExecute(subgraphName, executionRequest) {
@@ -137,8 +137,7 @@ export function unifiedGraphHandler(
               return opts.onSubgraphExecute(subgraphName, executionRequest);
             },
             queryPlan,
-          });
-        },
+          }),
       );
     },
   };

packages/router-runtime/src/types.ts

@@ -0,0 +1,26 @@
+import { QueryPlan } from '@graphql-hive/router-query-planner';
+import { ExecutionRequest, MaybePromise } from '@graphql-tools/utils';
+
+export type OnQueryPlanHook<TContext> = (
+  payload: OnQueryPlanHookPayload<TContext>,
+) => MaybePromise<OnQueryPlanDoneHook | void>;
+
+export type QueryPlanFn = (
+  executionRequest: ExecutionRequest,
+) => MaybePromise<QueryPlan>;
+
+export interface OnQueryPlanHookPayload<TContext> {
+  queryPlanFn: QueryPlanFn;
+  setQueryPlanFn(newQueryPlanFn: QueryPlanFn): void;
+  endQueryPlan(queryPlan: MaybePromise<QueryPlan>): void;
+  executionRequest: ExecutionRequest<any, TContext>;
+}
+
+export type OnQueryPlanDoneHook = (
+  payload: OnQueryPlanDoneHookPayload,
+) => MaybePromise<void>;
+
+export interface OnQueryPlanDoneHookPayload {
+  queryPlan: QueryPlan;
+  setQueryPlan(queryPlan: QueryPlan): void;
+}

packages/router-runtime/src/useQueryPlan.ts

@@ -1,18 +1,27 @@
 import type { GatewayPlugin } from '@graphql-hive/gateway-runtime';
 import type { QueryPlan } from '@graphql-hive/router-query-planner';
 import { isAsyncIterable } from '@graphql-tools/utils';
-import { queryPlanForExecutionRequestContext } from './utils';
 
 export interface QueryPlanOptions {
   /** Callback when the query plan has been successfuly generated. */
   onQueryPlan?(queryPlan: QueryPlan): void;
   /** Exposing the query plan inside the GraphQL result extensions. */
-  expose?: boolean | ((request: Request) => boolean);
+  exposeInResultExtensions?: boolean | ((request: Request) => boolean);
 }
 
 export function useQueryPlan(opts: QueryPlanOptions = {}): GatewayPlugin {
-  const { expose, onQueryPlan } = opts;
+  const queryPlanForExecutionRequestContext = new WeakMap<any, QueryPlan>();
+  const { exposeInResultExtensions, onQueryPlan } = opts;
   return {
+    onQueryPlan({ executionRequest }) {
+      return function onQueryPlanDone({ queryPlan }) {
+        queryPlanForExecutionRequestContext.set(
+          // getter like setter
+          executionRequest.context || executionRequest.document,
+          queryPlan,
+        );
+      };
+    },
     onExecute({ context, args }) {
       return {
         onExecuteDone({ result, setResult }) {
@@ -22,7 +31,9 @@ export function useQueryPlan(opts: QueryPlanOptions = {}): GatewayPlugin {
           );
           onQueryPlan?.(queryPlan!);
           const shouldExpose =
-            typeof expose === 'function' ? expose(context.request) : expose;
+            typeof exposeInResultExtensions === 'function'
+              ? exposeInResultExtensions(context.request)
+              : exposeInResultExtensions;
           if (shouldExpose && !isAsyncIterable(result)) {
             setResult({
               ...result,

packages/router-runtime/src/utils.ts

@@ -1,14 +1,8 @@
-import type { QueryPlan } from '@graphql-hive/router-query-planner';
 import {
   handleMaybePromise,
   type MaybePromise,
 } from '@whatwg-node/promise-helpers';
 
-export const queryPlanForExecutionRequestContext = new WeakMap<
-  any,
-  QueryPlan
->();
-
 export function getLazyPromise<T>(
   factory: () => MaybePromise<T>,
 ): () => MaybePromise<T> {

packages/router-runtime/src/wrapQueryPlanFnWithHooks.ts

@@ -0,0 +1,61 @@
+import { QueryPlan } from '@graphql-hive/router-query-planner';
+import { ExecutionRequest } from '@graphql-tools/utils';
+import {
+  handleMaybePromise,
+  iterateAsync,
+  MaybePromise,
+} from '@whatwg-node/promise-helpers';
+import { OnQueryPlanDoneHook, OnQueryPlanHook, QueryPlanFn } from './types';
+
+export function wrapQueryPlanFnWithHooks<TContext>(
+  defaultQueryPlanFn: QueryPlanFn,
+  onQueryPlanHooks: OnQueryPlanHook<TContext>[],
+): QueryPlanFn {
+  return function wrappedQueryPlanFn(
+    executionRequest: ExecutionRequest,
+  ): MaybePromise<QueryPlan> {
+    let queryPlanFn: QueryPlanFn = defaultQueryPlanFn;
+    let queryPlan$: MaybePromise<QueryPlan>;
+    const onQueryPlanDoneHooks: OnQueryPlanDoneHook[] = [];
+    return handleMaybePromise(
+      () =>
+        iterateAsync(
+          onQueryPlanHooks,
+          (onQueryPlanHook, endEarly) =>
+            onQueryPlanHook({
+              queryPlanFn,
+              setQueryPlanFn(newQueryPlanFn) {
+                queryPlanFn = newQueryPlanFn;
+              },
+              endQueryPlan(newQueryPlan) {
+                queryPlan$ = newQueryPlan;
+                endEarly();
+              },
+              executionRequest,
+            }),
+          onQueryPlanDoneHooks,
+        ),
+      () => {
+        if (queryPlan$) {
+          return queryPlan$;
+        }
+        return handleMaybePromise(
+          () => queryPlanFn(executionRequest),
+          (queryPlan) =>
+            handleMaybePromise(
+              () =>
+                iterateAsync(onQueryPlanDoneHooks, (onQueryPlanDoneHook) =>
+                  onQueryPlanDoneHook({
+                    queryPlan,
+                    setQueryPlan(newQueryPlan) {
+                      queryPlan = newQueryPlan;
+                    },
+                  }),
+                ),
+              () => queryPlan,
+            ),
+        );
+      },
+    );
+  };
+}

packages/router-runtime/tests/useQueryPlan.spec.ts

@@ -59,7 +59,7 @@ it('should include the query plan in result extensions when exposed', async () =
     unifiedGraphHandler,
     plugins: () => [
       useQueryPlan({
-        expose: true,
+        exposeInResultExtensions: true,
       }),
     ],
     subgraphs: [
@@ -114,7 +114,8 @@ it('should include the query plan in result extensions when expose returns true'
     unifiedGraphHandler,
     plugins: () => [
       useQueryPlan({
-        expose: (req) => req.headers.get('x-expose-query-plan') === 'true',
+        exposeInResultExtensions: (req) =>
+          req.headers.get('x-expose-query-plan') === 'true',
       }),
     ],
     subgraphs: [

packages/runtime/src/createGatewayRuntime.ts

@@ -82,6 +82,7 @@ import {
   type LandingPageRenderer,
   type YogaServerInstance,
 } from 'graphql-yoga';
+import { OnQueryPlanHook } from '../../router-runtime/src/types';
 import { createLoggerFromLogging } from './createLoggerFromLogging';
 import { createGraphOSFetcher } from './fetchers/graphos';
 import { getProxyExecutor } from './getProxyExecutor';
@@ -104,6 +105,7 @@ import { useDemandControl } from './plugins/useDemandControl';
 import { useFetchDebug } from './plugins/useFetchDebug';
 import useHiveConsole from './plugins/useHiveConsole';
 import { usePropagateHeaders } from './plugins/usePropagateHeaders';
+import { useMaybeQueryPlanDebug } from './plugins/useQueryPlanDebug';
 import { useRequestId } from './plugins/useRequestId';
 import { useRetryOnSchemaReload } from './plugins/useRetryOnSchemaReload';
 import { useSubgraphErrorPlugin } from './plugins/useSubgraphErrorPlugin';
@@ -187,6 +189,8 @@ export function createGatewayRuntime<
   const onDelegationStageExecuteHooks: OnDelegationStageExecuteHook<GatewayContext>[] =
     [];
 
+  const onQueryPlanHooks: OnQueryPlanHook<GatewayContext>[] = [];
+
   let unifiedGraph: GraphQLSchema;
   let schemaInvalidator: () => void;
   let getSchema: () => MaybePromise<GraphQLSchema> = () => unifiedGraph;
@@ -724,6 +728,7 @@ export function createGatewayRuntime<
       onSubgraphExecuteHooks,
       onDelegationPlanHooks,
       onDelegationStageExecuteHooks,
+      onQueryPlanHooks,
       additionalTypeDefs: config.additionalTypeDefs,
       additionalResolvers: config.additionalResolvers as IResolvers[],
       instrumentation: () => instrumentation,
@@ -817,6 +822,9 @@ export function createGatewayRuntime<
         if (plugin.onCacheDelete) {
           onCacheDeleteHooks.push(plugin.onCacheDelete);
         }
+        if (plugin.onQueryPlan) {
+          onQueryPlanHooks.push(plugin.onQueryPlan);
+        }
       }
     },
   };
@@ -1090,6 +1098,7 @@ export function createGatewayRuntime<
     useFetchDebug(),
     useMaybeDelegationPlanDebug({ log: configContext.log }),
     useCacheDebug({ log: configContext.log }),
+    useMaybeQueryPlanDebug({ log: configContext.log }),
   );
 
   const yoga = createYoga({