refactor(batchDelegate): simplify batch delegation

packages/batch-delegate/src/batchDelegateToSchema.ts

@@ -0,0 +1,10 @@
+import { BatchDelegateOptions, DataLoaderCache } from './types';
+
+import { getLoader } from './getLoader';
+
+export function batchDelegateToSchema<K = any, V = any, C = K>(options: BatchDelegateOptions): any {
+  let cache: DataLoaderCache<K, V, C>;
+
+  const loader = getLoader(cache, options);
+  return loader.load(options.key);
+}

packages/batch-delegate/src/createBatchDelegateFn.ts

@@ -1,53 +1,40 @@
-import { FieldNode, getNamedType, GraphQLOutputType, GraphQLList } from 'graphql';
-
 import DataLoader from 'dataloader';
 
-import { delegateToSchema } from '@graphql-tools/delegate';
+import {
+  CreateBatchDelegateFnOptions,
+  BatchDelegateOptionsFn,
+  BatchDelegateFn,
+  BatchDelegateArgsFn,
+  BatchDelegateResultsFn,
+  DataLoaderCache,
+} from './types';
 
-import { BatchDelegateOptionsFn, BatchDelegateFn, BatchDelegateOptions } from './types';
+import { getLoader } from './getLoader';
 
 export function createBatchDelegateFn<K = any, V = any, C = K>(
-  argFn: (args: ReadonlyArray<K>) => Record<string, any>,
-  batchDelegateOptionsFn: BatchDelegateOptionsFn,
+  optionsOrArgsFn: CreateBatchDelegateFnOptions | BatchDelegateArgsFn,
+  optionsFn?: BatchDelegateOptionsFn,
   dataLoaderOptions?: DataLoader.Options<K, V, C>,
-  resultsFn?: (results: any, keys: ReadonlyArray<K>) => V[]
+  resultsFn?: BatchDelegateResultsFn
 ): BatchDelegateFn<K> {
-  let cache: WeakMap<ReadonlyArray<FieldNode>, DataLoader<K, V, C>>;
-
-  function createBatchFn(options: BatchDelegateOptions) {
-    return async (keys: ReadonlyArray<K>) => {
-      let results = await delegateToSchema({
-        returnType: new GraphQLList(getNamedType(options.info.returnType) as GraphQLOutputType),
-        args: argFn(keys),
-        ...batchDelegateOptionsFn(options),
-      });
-      results = resultsFn ? resultsFn(results, keys) : results;
-      return Array.isArray(results) ? results : keys.map(() => results);
-    };
-  }
-
-  function getLoader(options: BatchDelegateOptions) {
-    if (!cache) {
-      cache = new WeakMap();
-      const batchFn = createBatchFn(options);
-      const newValue = new DataLoader<K, V, C>(keys => batchFn(keys), dataLoaderOptions);
-      cache.set(options.info.fieldNodes, newValue);
-      return newValue;
-    }
-
-    const cachedValue = cache.get(options.info.fieldNodes);
-    if (cachedValue === undefined) {
-      const batchFn = createBatchFn(options);
-      const newValue = new DataLoader<K, V, C>(keys => batchFn(keys), dataLoaderOptions);
-      cache.set(options.info.fieldNodes, newValue);
-      return newValue;
-    }
+  return typeof optionsOrArgsFn === 'function'
+    ? createBatchDelegateFnImpl({
+        argsFn: optionsOrArgsFn,
+        optionsFn,
+        dataLoaderOptions,
+        resultsFn,
+      })
+    : createBatchDelegateFnImpl(optionsOrArgsFn);
+}
 
-    return cachedValue;
-  }
+function createBatchDelegateFnImpl<K = any, V = any, C = K>(options: CreateBatchDelegateFnOptions): BatchDelegateFn<K> {
+  let cache: DataLoaderCache<K, V, C>;
 
-  return options => {
-    const loader = getLoader(options);
-    return loader.load(options.key);
+  return batchDelegateOptions => {
+    const loader = getLoader(cache, {
+      ...options,
+      ...batchDelegateOptions,
+    });
+    return loader.load(batchDelegateOptions.key);
   };
 }

packages/batch-delegate/src/getLoader.ts

@@ -0,0 +1,53 @@
+import { getNamedType, GraphQLOutputType, GraphQLList } from 'graphql';
+
+import DataLoader from 'dataloader';
+
+import { delegateToSchema } from '@graphql-tools/delegate';
+
+import { BatchDelegateOptions, DataLoaderCache } from './types';
+
+function createBatchFn<K = any>(options: BatchDelegateOptions) {
+  const argsFn = options.argsFn ?? ((keys: ReadonlyArray<K>) => ({ ids: keys }));
+  const { resultsFn, optionsFn } = options;
+
+  return async (keys: ReadonlyArray<K>) => {
+    let results = await delegateToSchema({
+      returnType: new GraphQLList(getNamedType(options.info.returnType) as GraphQLOutputType),
+      args: argsFn(keys),
+      ...(optionsFn != null ? optionsFn(options) : options),
+    });
+
+    if (resultsFn != null) {
+      results = resultsFn(results, keys);
+    }
+
+    return Array.isArray(results) ? results : keys.map(() => results);
+  };
+}
+
+function createLoader<K = any, V = any, C = K>(
+  cache: DataLoaderCache,
+  options: BatchDelegateOptions
+): DataLoader<K, V, C> {
+  const batchFn = createBatchFn(options);
+  const newValue = new DataLoader<K, V, C>(keys => batchFn(keys), options.dataLoaderOptions);
+  cache.set(options.info.fieldNodes, newValue);
+  return newValue;
+}
+
+export function getLoader<K = any, V = any, C = K>(
+  cache: DataLoaderCache,
+  options: BatchDelegateOptions
+): DataLoader<K, V, C> {
+  if (!cache) {
+    cache = new WeakMap();
+    return createLoader(cache, options);
+  }
+
+  const cachedValue = cache.get(options.info.fieldNodes);
+  if (cachedValue === undefined) {
+    return createLoader(cache, options);
+  }
+
+  return cachedValue;
+}

packages/batch-delegate/src/index.ts

@@ -1,3 +1,4 @@
-export { createBatchDelegateFn } from './createBatchDelegateFn';
+export * from './batchDelegateToSchema';
+export * from './createBatchDelegateFn';
 
 export * from './types';

packages/batch-delegate/src/types.ts

@@ -1,5 +1,11 @@
+import { FieldNode } from 'graphql';
+
+import DataLoader from 'dataloader';
+
 import { IDelegateToSchemaOptions } from '@graphql-tools/delegate';
 
+export type DataLoaderCache<K = any, V = any, C = K> = WeakMap<ReadonlyArray<FieldNode>, DataLoader<K, V, C>>;
+
 export type BatchDelegateFn<TContext = Record<string, any>, K = any> = (
   batchDelegateOptions: BatchDelegateOptions<TContext, K>
 ) => any;
@@ -8,7 +14,31 @@ export type BatchDelegateOptionsFn<TContext = Record<string, any>, K = any> = (
   batchDelegateOptions: BatchDelegateOptions<TContext, K>
 ) => IDelegateToSchemaOptions<TContext>;
 
-export interface BatchDelegateOptions<TContext = Record<string, any>, K = any>
+export type BatchDelegateArgsFn<K = any> = (keys: ReadonlyArray<K>) => Record<string, any>;
+
+export type BatchDelegateResultsFn<K = any, V = any> = (results: any, keys: ReadonlyArray<K>) => Array<V>;
+
+export interface BatchDelegateOptions<TContext = Record<string, any>, K = any, V = any, C = K>
   extends Omit<IDelegateToSchemaOptions<TContext>, 'args'> {
   key: K;
+  argsFn?: BatchDelegateArgsFn;
+  optionsFn?: BatchDelegateOptionsFn;
+  dataLoaderOptions?: DataLoader.Options<K, V, C>;
+  resultsFn?: BatchDelegateResultsFn;
+}
+
+export interface CreateBatchDelegateFnOptions<TContext = Record<string, any>, K = any, V = any, C = K>
+  extends Partial<Omit<IDelegateToSchemaOptions<TContext>, 'args' | 'info'>> {
+  dataLoaderOptions?: DataLoader.Options<K, V, C>;
+  argsFn?: BatchDelegateArgsFn;
+  resultsFn?: BatchDelegateResultsFn;
+  optionsFn?: BatchDelegateOptionsFn;
+}
+
+export interface BatchDelegateToSchema<TContext = Record<string, any>, K = any, V = any, C = K>
+  extends Omit<IDelegateToSchemaOptions<TContext>, 'schema' | 'args'> {
+  dataLoaderOptions?: DataLoader.Options<K, V, C>;
+  argsFn?: BatchDelegateArgsFn;
+  resultsFn?: BatchDelegateResultsFn;
+  optionsFn?: BatchDelegateOptionsFn;
 }

packages/delegate/src/types.ts

@@ -131,12 +131,14 @@ export interface SubschemaConfig {
   merge?: Record<string, MergedTypeConfig>;
 }
 
-export interface MergedTypeConfig {
+export interface MergedTypeConfig<K = any, V = any> {
   selectionSet?: string;
+  resolve?: MergedTypeResolver;
   fieldName?: string;
-  args?: (source: any) => Record<string, any>;
+  args?: (originalResult: any) => Record<string, any>;
   key?: (originalResult: any) => any;
-  resolve?: MergedTypeResolver;
+  argsFn?: (keys: ReadonlyArray<K>) => Record<string, any>;
+  resultsFn?: (results: any, keys: ReadonlyArray<K>) => Array<V>;
 }
 
 export type MergedTypeResolver = (

packages/stitch/src/stitchingInfo.ts

@@ -21,7 +21,7 @@ import {
 } from '@graphql-tools/utils';
 
 import { delegateToSchema, isSubschemaConfig, SubschemaConfig } from '@graphql-tools/delegate';
-import { createBatchDelegateFn } from '@graphql-tools/batch-delegate';
+import { batchDelegateToSchema } from '@graphql-tools/batch-delegate';
 
 import { MergeTypeCandidate, MergedTypeInfo, StitchingInfo, MergeTypeFilter } from './types';
 
@@ -120,26 +120,18 @@ function createMergedTypes(
           if (mergedTypeConfig.resolve != null) {
             targetSubschemas.push(subschema);
           } else if (mergedTypeConfig.key != null) {
-            const batchDelegateToSubschema = createBatchDelegateFn(
-              mergedTypeConfig.args,
-              ({ schema, selectionSet, context, info }) => ({
-                schema,
+            mergedTypeConfig.resolve = (originalResult, context, info, subschema, selectionSet) =>
+              batchDelegateToSchema({
+                schema: subschema,
                 operation: 'query',
                 fieldName: mergedTypeConfig.fieldName,
+                argsFn: mergedTypeConfig.argsFn ?? mergedTypeConfig.args,
+                resultsFn: mergedTypeConfig.resultsFn,
+                key: mergedTypeConfig.key(originalResult),
                 selectionSet,
                 context,
                 info,
                 skipTypeMerging: true,
-              })
-            );
-
-            mergedTypeConfig.resolve = (originalResult, context, info, subschema, selectionSet) =>
-              batchDelegateToSubschema({
-                key: mergedTypeConfig.key(originalResult),
-                schema: subschema,
-                context,
-                info,
-                selectionSet,
               });
 
             targetSubschemas.push(subschema);