From 70d42cf8d4323efe55a5d0280687d9b48b8c4856 Mon Sep 17 00:00:00 2001
From: Benjie Gillam <benjie@jemjie.com>
Date: Mon, 20 May 2024 18:21:55 +0100
Subject: [PATCH] Middleware around establishing an operation plan

---
 grafast/grafast/src/index.ts      |  2 ++
 grafast/grafast/src/interfaces.ts | 12 +++++++
 grafast/grafast/src/prepare.ts    | 57 +++++++++++++++++++++++++------
 3 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/grafast/grafast/src/index.ts b/grafast/grafast/src/index.ts
index 9458f9aec9..70ad2c2deb 100644
--- a/grafast/grafast/src/index.ts
+++ b/grafast/grafast/src/index.ts
@@ -45,6 +45,7 @@ import type {
   $$queryCache,
   CacheByOperationEntry,
   DataFromStep,
+  EstablishOperationPlanEvent,
   ExecuteEvent,
   GrafastExecutionArgs,
   GrafastTimeouts,
@@ -756,6 +757,7 @@ declare global {
       ): PromiseOrDirect<Grafast.ExecutionArgs>;
       execute(event: ExecuteEvent): ReturnType<typeof execute>;
       subscribe(event: ExecuteEvent): ReturnType<typeof subscribe>;
+      establishOperationPlan(event: EstablishOperationPlanEvent): OperationPlan;
     }
     interface Plugin {
       grafast?: {
diff --git a/grafast/grafast/src/interfaces.ts b/grafast/grafast/src/interfaces.ts
index bd7c5341be..5117f9fd01 100644
--- a/grafast/grafast/src/interfaces.ts
+++ b/grafast/grafast/src/interfaces.ts
@@ -19,10 +19,12 @@ import type {
   GraphQLScalarType,
   GraphQLSchema,
   GraphQLType,
+  OperationDefinitionNode,
   Source,
   ValueNode,
   VariableNode,
 } from "graphql";
+import type { ObjMap } from "graphql/jsutils/ObjMap.js";
 
 import type { Bucket, RequestTools } from "./bucket.js";
 import type { OperationPlan } from "./engine/OperationPlan.js";
@@ -972,3 +974,13 @@ export interface ExecuteEvent {
 export interface SubscribeEvent {
   args: GrafastExecutionArgs;
 }
+export interface EstablishOperationPlanEvent {
+  schema: GraphQLSchema;
+  operation: OperationDefinitionNode;
+  fragments: ObjMap<FragmentDefinitionNode>;
+  variableValues: Record<string, any>;
+  context: any;
+  rootValue: any;
+  planningTimeout: number | undefined;
+  args: GrafastExecutionArgs;
+}
diff --git a/grafast/grafast/src/prepare.ts b/grafast/grafast/src/prepare.ts
index 161b83226e..fb3c8deace 100644
--- a/grafast/grafast/src/prepare.ts
+++ b/grafast/grafast/src/prepare.ts
@@ -29,8 +29,13 @@ import { POLYMORPHIC_ROOT_PATH } from "./engine/OperationPlan.js";
 import type { OutputPlan } from "./engine/OutputPlan.js";
 import { coerceError, getChildBucketAndIndex } from "./engine/OutputPlan.js";
 import { establishOperationPlan } from "./establishOperationPlan.js";
-import type { GrafastPlanJSON, OperationPlan } from "./index.js";
 import type {
+  GrafastExecutionArgs,
+  GrafastPlanJSON,
+  OperationPlan,
+} from "./index.js";
+import type {
+  EstablishOperationPlanEvent,
   GrafastTimeouts,
   JSONValue,
   PromiseOrDirect,
@@ -537,11 +542,23 @@ declare module "./engine/OperationPlan.js" {
   }
 }
 
+function establishOperationPlanFromEvent(event: EstablishOperationPlanEvent) {
+  return establishOperationPlan(
+    event.schema,
+    event.operation,
+    event.fragments,
+    event.variableValues,
+    event.context as any,
+    event.rootValue,
+    event.planningTimeout,
+  );
+}
+
 /**
  * @internal
  */
 export function grafastPrepare(
-  args: ExecutionArgs,
+  args: GrafastExecutionArgs,
   options: GrafastPrepareOptions = {},
 ): PromiseOrDirect<
   ExecutionResult | AsyncGenerator<AsyncExecutionResult, void, void>
@@ -552,6 +569,7 @@ export function grafastPrepare(
     rootValue = Object.create(null),
     // operationName,
     // document,
+    middlewares,
   } = args;
   const exeContext = buildExecutionContext(args);
 
@@ -567,15 +585,32 @@ export function grafastPrepare(
   const planningTimeout = options.timeouts?.planning;
   let operationPlan!: OperationPlan;
   try {
-    operationPlan = establishOperationPlan(
-      schema,
-      operation,
-      fragments,
-      variableValues,
-      context as any,
-      rootValue,
-      planningTimeout,
-    );
+    if (middlewares) {
+      operationPlan = middlewares.run(
+        "establishOperationPlan",
+        {
+          schema,
+          operation,
+          fragments,
+          variableValues,
+          context: context as any,
+          rootValue,
+          planningTimeout,
+          args,
+        },
+        establishOperationPlanFromEvent,
+      );
+    } else {
+      operationPlan = establishOperationPlan(
+        schema,
+        operation,
+        fragments,
+        variableValues,
+        context as any,
+        rootValue,
+        planningTimeout,
+      );
+    }
   } catch (error) {
     const graphqlError =
       error instanceof GraphQLError