Expose the query plan in graphql extensions for Rust QP

.changeset/violet-jokes-enjoy.md

@@ -0,0 +1,5 @@
+---
+'@graphql-hive/router-runtime': patch
+---
+
+Expose the query plan in graphql result extensions when setting `queryPlanInExtensions` to true in context

packages/router-runtime/src/handler.ts

@@ -84,51 +84,66 @@ export function unifiedGraphHandler(
             executionRequest.operationName || null,
           ),
         (queryPlan) =>
-          executeQueryPlan({
-            supergraphSchema,
-            executionRequest,
-            onSubgraphExecute(subgraphName, executionRequest) {
-              const subschema = getSubschema(subgraphName);
-              if (subschema.transforms?.length) {
-                const transforms = subschema.transforms;
-                const transformationContext = Object.create(null);
-                for (const transform of transforms) {
-                  if (transform.transformRequest) {
-                    executionRequest = transform.transformRequest(
-                      executionRequest,
-                      undefined as any,
-                      transformationContext,
-                    );
-                  }
-                }
-                return handleMaybePromise(
-                  () => opts.onSubgraphExecute(subgraphName, executionRequest),
-                  (executionResult) => {
-                    function handleResult(executionResult: ExecutionResult) {
-                      for (const transform of transforms.toReversed()) {
-                        if (transform.transformResult) {
-                          executionResult = transform.transformResult(
-                            executionResult,
-                            undefined as any,
-                            transformationContext,
-                          );
-                        }
+          handleMaybePromise(
+            () =>
+              executeQueryPlan({
+                supergraphSchema,
+                executionRequest,
+                onSubgraphExecute(subgraphName, executionRequest) {
+                  const subschema = getSubschema(subgraphName);
+                  if (subschema.transforms?.length) {
+                    const transforms = subschema.transforms;
+                    const transformationContext = Object.create(null);
+                    for (const transform of transforms) {
+                      if (transform.transformRequest) {
+                        executionRequest = transform.transformRequest(
+                          executionRequest,
+                          undefined as any,
+                          transformationContext,
+                        );
                       }
-                      return executionResult;
-                    }
-                    if (isAsyncIterable(executionResult)) {
-                      return mapAsyncIterator(executionResult, (result) =>
-                        handleResult(result),
-                      );
                     }
-                    return handleResult(executionResult);
-                  },
-                );
+                    return handleMaybePromise(
+                      () =>
+                        opts.onSubgraphExecute(subgraphName, executionRequest),
+                      (executionResult) => {
+                        function handleResult(
+                          executionResult: ExecutionResult,
+                        ) {
+                          for (const transform of transforms.toReversed()) {
+                            if (transform.transformResult) {
+                              executionResult = transform.transformResult(
+                                executionResult,
+                                undefined as any,
+                                transformationContext,
+                              );
+                            }
+                          }
+                          return executionResult;
+                        }
+                        if (isAsyncIterable(executionResult)) {
+                          return mapAsyncIterator(executionResult, (result) =>
+                            handleResult(result),
+                          );
+                        }
+                        return handleResult(executionResult);
+                      },
+                    );
+                  }
+                  return opts.onSubgraphExecute(subgraphName, executionRequest);
+                },
+                queryPlan,
+              }),
+            (result) => {
+              if (isAsyncIterable(result)) {
+                return result;
+              }
+              if (executionRequest.context?.['queryPlanInExtensions']) {
+                result.extensions = { queryPlan };
               }
-              return opts.onSubgraphExecute(subgraphName, executionRequest);
+              return result;
             },
-            queryPlan,
-          }),
+          ),
       );
     },
   };

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

@@ -0,0 +1,84 @@
+import { GatewayPlugin } from '@graphql-hive/gateway-runtime';
+import { createGatewayTester } from '@graphql-hive/gateway-testing';
+import { expect, it } from 'vitest';
+import { unifiedGraphHandler } from '../src/handler';
+
+it('should include the query plan in result extensions when context', async () => {
+  await using gw = createGatewayTester({
+    unifiedGraphHandler,
+    plugins: () => [
+      {
+        onContextBuilding({ context, extendContext }) {
+          if (context.headers['hive-expose-query-plan'] === 'true') {
+            extendContext({
+              queryPlanInExtensions: true,
+            });
+          }
+        },
+      } satisfies GatewayPlugin,
+    ],
+    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({
+      query: /* GraphQL */ `
+        {
+          hello
+        }
+      `,
+      headers: {
+        'hive-expose-query-plan': 'true',
+      },
+    }),
+  ).resolves.toMatchInlineSnapshot(`
+    {
+      "data": {
+        "hello": "world",
+      },
+      "extensions": {
+        "queryPlan": {
+          "kind": "QueryPlan",
+          "node": {
+            "kind": "Fetch",
+            "operation": "{hello}",
+            "operationKind": "query",
+            "serviceName": "upstream",
+          },
+        },
+      },
+    }
+  `);
+});