diff --git a/.changeset/four-ads-yell.md b/.changeset/four-ads-yell.md new file mode 100644 index 000000000..69a1654fe --- /dev/null +++ b/.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 diff --git a/packages/fusion-runtime/src/unifiedGraphManager.ts b/packages/fusion-runtime/src/unifiedGraphManager.ts index 7fdba021b..ea27c3f70 100644 --- a/packages/fusion-runtime/src/unifiedGraphManager.ts +++ b/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[]; onDelegationStageExecuteHooks?: OnDelegationStageExecuteHook[]; onDelegateHooks?: OnDelegateHook[]; + onQueryPlanHooks?: OnQueryPlanHook[]; /** * Configure the batch delegation options for all merged types in all subschemas. */ @@ -119,6 +121,7 @@ export interface UnifiedGraphManagerOptions { onDelegateHooks?: OnDelegateHook[]; onDelegationPlanHooks?: OnDelegationPlanHook[]; onDelegationStageExecuteHooks?: OnDelegationStageExecuteHook[]; + onQueryPlanHooks?: OnQueryPlanHook[]; /** * Whether to batch the subgraph executions. * @default true @@ -167,6 +170,7 @@ export class UnifiedGraphManager implements AsyncDisposable { private onSubgraphExecuteHooks: OnSubgraphExecuteHook[]; private onDelegationPlanHooks: OnDelegationPlanHook[]; private onDelegationStageExecuteHooks: OnDelegationStageExecuteHook[]; + private onQueryPlanHooks: OnQueryPlanHook[]; private inContextSDK: any; private initialUnifiedGraph$?: MaybePromise; private polling$?: MaybePromise; @@ -189,6 +193,7 @@ export class UnifiedGraphManager 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 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 diff --git a/packages/fusion-runtime/src/utils.ts b/packages/fusion-runtime/src/utils.ts index 64d9d45c9..07bf15928 100644 --- a/packages/fusion-runtime/src/utils.ts +++ b/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 { onSubgraphExecute?: OnSubgraphExecuteHook; onDelegationPlan?: OnDelegationPlanHook; onDelegationStageExecute?: OnDelegationStageExecuteHook; + onQueryPlan?: OnQueryPlanHook; } export type OnSubgraphExecuteHook = ( diff --git a/packages/router-runtime/src/handler.ts b/packages/router-runtime/src/handler.ts index 98c212d68..4969aa609 100644 --- a/packages/router-runtime/src/handler.ts +++ b/packages/router-runtime/src/handler.ts @@ -7,7 +7,11 @@ 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, @@ -15,11 +19,8 @@ import { } 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> >(); - function planDocument(document: DocumentNode, operationName: string | null) { + function defaultQueryPlanFn({ + document, + operationName, + }: ExecutionRequest): MaybePromise { 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, - }); - }, + }), ); }, }; diff --git a/packages/router-runtime/src/types.ts b/packages/router-runtime/src/types.ts new file mode 100644 index 000000000..c7661d30d --- /dev/null +++ b/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 = ( + payload: OnQueryPlanHookPayload, +) => MaybePromise; + +export type QueryPlanFn = ( + executionRequest: ExecutionRequest, +) => MaybePromise; + +export interface OnQueryPlanHookPayload { + queryPlanFn: QueryPlanFn; + setQueryPlanFn(newQueryPlanFn: QueryPlanFn): void; + endQueryPlan(queryPlan: MaybePromise): void; + executionRequest: ExecutionRequest; +} + +export type OnQueryPlanDoneHook = ( + payload: OnQueryPlanDoneHookPayload, +) => MaybePromise; + +export interface OnQueryPlanDoneHookPayload { + queryPlan: QueryPlan; + setQueryPlan(queryPlan: QueryPlan): void; +} diff --git a/packages/router-runtime/src/useQueryPlan.ts b/packages/router-runtime/src/useQueryPlan.ts index 5160cc6f6..6eeff8fc8 100644 --- a/packages/router-runtime/src/useQueryPlan.ts +++ b/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(); + 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, diff --git a/packages/router-runtime/src/utils.ts b/packages/router-runtime/src/utils.ts index d5a91f77e..a8b4326c1 100644 --- a/packages/router-runtime/src/utils.ts +++ b/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( factory: () => MaybePromise, ): () => MaybePromise { diff --git a/packages/router-runtime/src/wrapQueryPlanFnWithHooks.ts b/packages/router-runtime/src/wrapQueryPlanFnWithHooks.ts new file mode 100644 index 000000000..9066278c0 --- /dev/null +++ b/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( + defaultQueryPlanFn: QueryPlanFn, + onQueryPlanHooks: OnQueryPlanHook[], +): QueryPlanFn { + return function wrappedQueryPlanFn( + executionRequest: ExecutionRequest, + ): MaybePromise { + let queryPlanFn: QueryPlanFn = defaultQueryPlanFn; + let queryPlan$: MaybePromise; + 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, + ), + ); + }, + ); + }; +} diff --git a/packages/router-runtime/tests/useQueryPlan.spec.ts b/packages/router-runtime/tests/useQueryPlan.spec.ts index a85c6f282..400e74884 100644 --- a/packages/router-runtime/tests/useQueryPlan.spec.ts +++ b/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: [ diff --git a/packages/runtime/src/createGatewayRuntime.ts b/packages/runtime/src/createGatewayRuntime.ts index 48c5bf36e..e0d1c1c29 100644 --- a/packages/runtime/src/createGatewayRuntime.ts +++ b/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[] = []; + const onQueryPlanHooks: OnQueryPlanHook[] = []; + let unifiedGraph: GraphQLSchema; let schemaInvalidator: () => void; let getSchema: () => MaybePromise = () => 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({ diff --git a/packages/runtime/src/plugins/useQueryPlanDebug.ts b/packages/runtime/src/plugins/useQueryPlanDebug.ts new file mode 100644 index 000000000..f1ea5c613 --- /dev/null +++ b/packages/runtime/src/plugins/useQueryPlanDebug.ts @@ -0,0 +1,54 @@ +import { Logger } from '@graphql-hive/logger'; +import type { GatewayPlugin } from '../types'; + +export function useMaybeQueryPlanDebug>({ + log, +}: { + log: Logger; +}): GatewayPlugin { + let activePlugin: GatewayPlugin | undefined; + return { + onPluginInit({ plugins }) { + let shouldLog = false; + log.debug(() => (shouldLog = true)); + if (shouldLog) { + activePlugin = useQueryPlanDebug(); + // plugins.push will run the plugin last, but addPlugin will run it after this plugin. we dont care? + plugins.push( + // @ts-expect-error TODO: fix types + activePlugin, + ); + } else if (activePlugin) { + const index = plugins.indexOf( + // @ts-expect-error TODO: fix types + activePlugin, + ); + if ( + // must be + index > -1 + ) { + plugins.splice(index, 1); + } + activePlugin = undefined; + } + }, + }; +} + +function useQueryPlanDebug< + TContext extends Record, +>(): GatewayPlugin { + return { + onQueryPlan({ executionRequest }) { + return ({ queryPlan }) => { + executionRequest.context?.log.debug( + { + queryPlan, + operationName: executionRequest.operationName || 'Anonymous', + }, + `[useQueryPlanDebug] `, + ); + }; + }, + }; +}