From 0c9299866fe1e26e361cbf967916315f12b9581d Mon Sep 17 00:00:00 2001
From: Yaacov Rydzinski <yaacovCR@gmail.com>
Date: Mon, 31 Aug 2020 21:51:39 +0300
Subject: [PATCH] fix(delegate): pass extensions to executor/subscriber

---
 packages/delegate/src/delegateToSchema.ts     |  6 ++----
 packages/delegate/src/getBatchingExecutor.ts  | 21 ++++++++++++++++---
 packages/delegate/src/mergeExecutionParams.ts |  8 ++++++-
 packages/delegate/src/types.ts                |  5 +++++
 4 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/packages/delegate/src/delegateToSchema.ts b/packages/delegate/src/delegateToSchema.ts
index b91cdd4da93..7c86e17d6de 100644
--- a/packages/delegate/src/delegateToSchema.ts
+++ b/packages/delegate/src/delegateToSchema.ts
@@ -163,8 +163,7 @@ export function delegateRequest({
     }
 
     const executionResult = executor({
-      document: processedRequest.document,
-      variables: processedRequest.variables,
+      ...processedRequest,
       context,
       info,
     });
@@ -179,8 +178,7 @@ export function delegateRequest({
     subschemaConfig?.subscriber || createDefaultSubscriber(targetSchema, subschemaConfig?.rootValue || targetRootValue);
 
   return subscriber({
-    document: processedRequest.document,
-    variables: processedRequest.variables,
+    ...processedRequest,
     context,
     info,
   }).then((subscriptionResult: AsyncIterableIterator<ExecutionResult> | ExecutionResult) => {
diff --git a/packages/delegate/src/getBatchingExecutor.ts b/packages/delegate/src/getBatchingExecutor.ts
index 9ce59e56439..14f581c3deb 100644
--- a/packages/delegate/src/getBatchingExecutor.ts
+++ b/packages/delegate/src/getBatchingExecutor.ts
@@ -17,14 +17,18 @@ export const getBatchingExecutor = memoize2of3(function (
   executor: ({ document, context, variables, info }: ExecutionParams) => ExecutionResult | Promise<ExecutionResult>
 ) {
   const loader = new DataLoader(
-    createLoadFn(executor ?? subschemaConfig.executor),
+    createLoadFn(
+      executor ?? subschemaConfig.executor,
+      subschemaConfig.batchingOptions?.extensionsReducer ?? defaultExtensionsReducer
+    ),
     subschemaConfig.batchingOptions?.dataLoaderOptions
   );
   return (executionParams: ExecutionParams) => loader.load(executionParams);
 });
 
 function createLoadFn(
-  executor: ({ document, context, variables, info }: ExecutionParams) => ExecutionResult | Promise<ExecutionResult>
+  executor: ({ document, context, variables, info }: ExecutionParams) => ExecutionResult | Promise<ExecutionResult>,
+  extensionsReducer: (mergedExtensions: Record<string, any>, executionParams: ExecutionParams) => Record<string, any>
 ) {
   return async (execs: Array<ExecutionParams>): Promise<Array<ExecutionResult>> => {
     const execBatches: Array<Array<ExecutionParams>> = [];
@@ -46,7 +50,7 @@ function createLoadFn(
     let containsPromises = false;
     const executionResults: Array<ExecutionResult | Promise<ExecutionResult>> = [];
     execBatches.forEach(execBatch => {
-      const mergedExecutionParams = mergeExecutionParams(execBatch);
+      const mergedExecutionParams = mergeExecutionParams(execBatch, extensionsReducer);
       const executionResult = executor(mergedExecutionParams);
 
       if (isPromise(executionResult)) {
@@ -72,3 +76,14 @@ function createLoadFn(
     return results;
   };
 }
+
+function defaultExtensionsReducer(
+  mergedExtensions: Record<string, any>,
+  executionParams: ExecutionParams
+): Record<string, any> {
+  const newExtensions = executionParams.extensions;
+  if (newExtensions != null) {
+    Object.assign(mergedExtensions, newExtensions);
+  }
+  return mergedExtensions;
+}
diff --git a/packages/delegate/src/mergeExecutionParams.ts b/packages/delegate/src/mergeExecutionParams.ts
index e97c0e44028..9c8cfbc6af9 100644
--- a/packages/delegate/src/mergeExecutionParams.ts
+++ b/packages/delegate/src/mergeExecutionParams.ts
@@ -54,11 +54,15 @@ import { ExecutionParams } from './types';
  *     }
  *   }
  */
-export function mergeExecutionParams(execs: Array<ExecutionParams>): ExecutionParams {
+export function mergeExecutionParams(
+  execs: Array<ExecutionParams>,
+  extensionsReducer: (mergedExtensions: Record<string, any>, executionParams: ExecutionParams) => Record<string, any>
+): ExecutionParams {
   const mergedVariables: Record<string, any> = Object.create(null);
   const mergedVariableDefinitions: Array<VariableDefinitionNode> = [];
   const mergedSelections: Array<SelectionNode> = [];
   const mergedFragmentDefinitions: Array<FragmentDefinitionNode> = [];
+  let mergedExtensions: Record<string, any> = Object.create(null);
 
   execs.forEach((executionParams, index) => {
     const prefixedExecutionParams = prefixExecutionParams(createPrefix(index), executionParams);
@@ -73,6 +77,7 @@ export function mergeExecutionParams(execs: Array<ExecutionParams>): ExecutionPa
       }
     });
     Object.assign(mergedVariables, prefixedExecutionParams.variables);
+    mergedExtensions = extensionsReducer(mergedExtensions, executionParams);
   });
 
   const mergedOperationDefinition: OperationDefinitionNode = {
@@ -91,6 +96,7 @@ export function mergeExecutionParams(execs: Array<ExecutionParams>): ExecutionPa
       definitions: [mergedOperationDefinition, ...mergedFragmentDefinitions],
     },
     variables: mergedVariables,
+    extensions: mergedExtensions,
     context: execs[0].context,
     info: execs[0].info,
   };
diff --git a/packages/delegate/src/types.ts b/packages/delegate/src/types.ts
index b4868f67962..3573506f40b 100644
--- a/packages/delegate/src/types.ts
+++ b/packages/delegate/src/types.ts
@@ -95,6 +95,7 @@ export interface MergedTypeInfo {
 export interface ExecutionParams<TArgs = Record<string, any>, TContext = any> {
   document: DocumentNode;
   variables?: TArgs;
+  extensions?: Record<string, any>;
   context?: TContext;
   info?: GraphQLResolveInfo;
 }
@@ -134,6 +135,10 @@ export interface SubschemaConfig<K = any, V = any, C = K> {
   merge?: Record<string, MergedTypeConfig>;
   batch?: boolean;
   batchingOptions?: {
+    extensionsReducer?: (
+      mergedExtensions: Record<string, any>,
+      executionParams: ExecutionParams
+    ) => Record<string, any>;
     dataLoaderOptions?: DataLoader.Options<K, V, C>;
   };
 }