Expose the query plan in graphql extensions for Rust QP (#1739)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: enisdenjo

.changeset/@graphql-hive_gateway-runtime-1739-dependencies.md

@@ -0,0 +1,7 @@
+---
+'@graphql-hive/gateway-runtime': patch
+---
+
+dependencies updates: 
+
+- Updated dependency [`@graphql-yoga/plugin-apollo-usage-report@^0.12.0` ↗︎](https://www.npmjs.com/package/@graphql-yoga/plugin-apollo-usage-report/v/0.12.0) (from `^0.11.2`, in `dependencies`)

.changeset/@graphql-hive_router-runtime-1739-dependencies.md

@@ -0,0 +1,7 @@
+---
+'@graphql-hive/router-runtime': patch
+---
+
+dependencies updates: 
+
+- Updated dependency [`@graphql-hive/router-query-planner@^0.0.6` ↗︎](https://www.npmjs.com/package/@graphql-hive/router-query-planner/v/0.0.6) (from `^0.0.4`, in `dependencies`)

.changeset/violet-jokes-enjoy.md

@@ -0,0 +1,5 @@
+---
+'@graphql-hive/router-runtime': patch
+---
+
+Expose the query plan by using the `useQueryPlan` plugin

packages/router-runtime/src/handler.ts

@@ -15,7 +15,11 @@ import {
 } from '@whatwg-node/promise-helpers';
 import { BREAK, DocumentNode, visit } from 'graphql';
 import { executeQueryPlan } from './executor';
-import { getLazyFactory, getLazyValue } from './utils';
+import {
+  getLazyFactory,
+  getLazyValue,
+  queryPlanForExecutionRequestContext,
+} from './utils';
 
 export function unifiedGraphHandler(
   opts: UnifiedGraphHandlerOpts,
@@ -83,8 +87,13 @@ export function unifiedGraphHandler(
             executionRequest.document,
             executionRequest.operationName || null,
           ),
-        (queryPlan) =>
-          executeQueryPlan({
+        (queryPlan) => {
+          queryPlanForExecutionRequestContext.set(
+            // setter like getter
+            executionRequest.context || executionRequest.document,
+            queryPlan,
+          );
+          return executeQueryPlan({
             supergraphSchema,
             executionRequest,
             onSubgraphExecute(subgraphName, executionRequest) {
@@ -128,7 +137,8 @@ export function unifiedGraphHandler(
               return opts.onSubgraphExecute(subgraphName, executionRequest);
             },
             queryPlan,
-          }),
+          });
+        },
       );
     },
   };

packages/router-runtime/src/index.ts

@@ -1 +1,2 @@
 export * from './handler';
+export * from './useQueryPlan';

packages/router-runtime/src/useQueryPlan.ts

@@ -0,0 +1,39 @@
+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);
+}
+
+export function useQueryPlan(opts: QueryPlanOptions = {}): GatewayPlugin {
+  const { expose, onQueryPlan } = opts;
+  return {
+    onExecute({ context, args }) {
+      return {
+        onExecuteDone({ result, setResult }) {
+          const queryPlan = queryPlanForExecutionRequestContext.get(
+            // getter like setter
+            context || args.document,
+          );
+          onQueryPlan?.(queryPlan!);
+          const shouldExpose =
+            typeof expose === 'function' ? expose(context.request) : expose;
+          if (shouldExpose && !isAsyncIterable(result)) {
+            setResult({
+              ...result,
+              extensions: {
+                ...result.extensions,
+                queryPlan,
+              },
+            });
+          }
+        },
+      };
+    },
+  };
+}

packages/router-runtime/src/utils.ts

@@ -1,8 +1,14 @@
+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/tests/useQueryPlan.spec.ts

@@ -0,0 +1,184 @@
+import { createGatewayTester } from '@graphql-hive/gateway-testing';
+import { QueryPlan } from '@graphql-hive/router-query-planner';
+import { expect, it } from 'vitest';
+import { unifiedGraphHandler, useQueryPlan } from '../src/index';
+
+it('should callback when the query plan is available', async () => {
+  let queryPlan!: QueryPlan;
+  await using gw = createGatewayTester({
+    unifiedGraphHandler,
+    plugins: () => [
+      useQueryPlan({
+        onQueryPlan(_queryPlan) {
+          queryPlan = _queryPlan;
+        },
+      }),
+    ],
+    subgraphs: [
+      {
+        name: 'upstream',
+        schema: {
+          typeDefs: /* GraphQL */ `
+            type Query {
+              hello: String
+            }
+          `,
+          resolvers: {
+            Query: {
+              hello: () => 'world',
+            },
+          },
+        },
+      },
+    ],
+  });
+
+  await gw.execute({
+    query: /* GraphQL */ `
+      {
+        hello
+      }
+    `,
+  });
+
+  expect(queryPlan).toMatchInlineSnapshot(`
+    {
+      "kind": "QueryPlan",
+      "node": {
+        "kind": "Fetch",
+        "operation": "{hello}",
+        "operationKind": "query",
+        "serviceName": "upstream",
+      },
+    }
+  `);
+});
+
+it('should include the query plan in result extensions when exposed', async () => {
+  await using gw = createGatewayTester({
+    unifiedGraphHandler,
+    plugins: () => [
+      useQueryPlan({
+        expose: true,
+      }),
+    ],
+    subgraphs: [
+      {
+        name: 'upstream',
+        schema: {
+          typeDefs: /* GraphQL */ `
+            type Query {
+              hello: String
+            }
+          `,
+          resolvers: {
+            Query: {
+              hello: () => 'world',
+            },
+          },
+        },
+      },
+    ],
+  });
+
+  await expect(
+    gw.execute({
+      query: /* GraphQL */ `
+        {
+          hello
+        }
+      `,
+    }),
+  ).resolves.toMatchInlineSnapshot(`
+    {
+      "data": {
+        "hello": "world",
+      },
+      "extensions": {
+        "queryPlan": {
+          "kind": "QueryPlan",
+          "node": {
+            "kind": "Fetch",
+            "operation": "{hello}",
+            "operationKind": "query",
+            "serviceName": "upstream",
+          },
+        },
+      },
+    }
+  `);
+});
+
+it('should include the query plan in result extensions when expose returns true', async () => {
+  await using gw = createGatewayTester({
+    unifiedGraphHandler,
+    plugins: () => [
+      useQueryPlan({
+        expose: (req) => req.headers.get('x-expose-query-plan') === 'true',
+      }),
+    ],
+    subgraphs: [
+      {
+        name: 'upstream',
+        schema: {
+          typeDefs: /* GraphQL */ `
+            type Query {
+              hello: String
+            }
+          `,
+          resolvers: {
+            Query: {
+              hello: () => 'world',
+            },
+          },
+        },
+      },
+    ],
+  });
+
+  await expect(
+    gw.execute({
+      query: /* GraphQL */ `
+        {
+          hello
+        }
+      `,
+    }),
+  ).resolves.toMatchInlineSnapshot(`
+    {
+      "data": {
+        "hello": "world",
+      },
+    }
+  `);
+
+  await expect(
+    gw.execute({
+      headers: {
+        'x-expose-query-plan': 'true',
+      },
+      query: /* GraphQL */ `
+        {
+          hello
+        }
+      `,
+    }),
+  ).resolves.toMatchInlineSnapshot(`
+    {
+      "data": {
+        "hello": "world",
+      },
+      "extensions": {
+        "queryPlan": {
+          "kind": "QueryPlan",
+          "node": {
+            "kind": "Fetch",
+            "operation": "{hello}",
+            "operationKind": "query",
+            "serviceName": "upstream",
+          },
+        },
+      },
+    }
+  `);
+});