enhance(stitching): increase key flexibility when type merging (#1888)

* demonstrate merging of federated-like schemas

* enhance(stitching): only add necessary keys

also consolidates equivalent selections

add memoization

This could be optimized differently: if there was a mapping of fields within the gateway schema to the subschemas containing those fields, we could instead iterate through the document and collect the originating subschemas and then add the keys for each of those subschemas. This method seems possibly about as good.

* refactor: rename variable to be more descriptive

* enhance(typeMerging): expand allowed keys

allow keys whose source data is spread across multiple schemas

* fix(selectionSets): only spread correct fragments

* refactor(typeMerging): reduce memoization arguments

* enhance(stitching): consolidate selections

* proof of concept?

* fix

* enhance(typeMerging): prune fields that cause key fields to be added

key fields for a given subschema by definition must be resolved outside that subschema, and so for any given key field, there is no need to add the other key fields to the selectionSet or to consider the key field as proxiable via that subschema.

The key field is the "key" to the remaining fields, but -- for the purposes of type merging -- will never be obtained from the given subschema.

This becomes more complex with non-leaf key fields, and so this pruning is enabled only for leaf key fields.

* fix

* change example

to favor input object type over scalar

see #1888 (review)

packages/batch-delegate/tests/typeMerging.example.test.ts

@@ -56,17 +56,18 @@ describe('merging using type merging', () => {
 
   const inventorySchema = makeExecutableSchema({
     typeDefs: `
+      input ProductRepresentation {
+        upc: String!
+        price: Int
+        weight: Int
+      }
       type Product {
         upc: String!
         inStock: Boolean
         shippingEstimate: Int
       }
       type Query {
-        _productByUpc(
-          upc: String!,
-          weight: Int,
-          price: Int
-        ): Product
+        _productByRepresentation(product: ProductRepresentation): Product
       }
     `,
     resolvers: {
@@ -79,10 +80,12 @@ describe('merging using type merging', () => {
         }
       },
       Query: {
-        _productByUpc: (_root, { upc, ...fields }) => ({
-          ...inventory.find(product => product.upc === upc),
-          ...fields
-        }),
+        _productByRepresentation: (_root, { product: { upc, ...fields } }) => {
+          return {
+            ...inventory.find(product => product.upc === upc),
+            ...fields
+          };
+        },
       },
     },
   });
@@ -230,9 +233,14 @@ describe('merging using type merging', () => {
         schema: inventorySchema,
         merge: {
           Product: {
-            selectionSet: '{ upc weight price }',
-            fieldName: '_productByUpc',
-            args: ({ upc, weight, price }) => ({ upc, weight, price }),
+            selectionSet: '{ upc }',
+            fields: {
+              shippingEstimate: {
+                selectionSet: '{ price weight }',
+              },
+            },
+            fieldName: '_productByRepresentation',
+            args: ({ upc, weight, price }) => ({ product: { upc, weight, price } }),
           }
         }
       },

packages/delegate/src/delegationBindings.ts

@@ -41,7 +41,7 @@ export function defaultDelegationBinding(delegationContext: DelegationContext):
       new AddSelectionSets(
         info.schema,
         returnType,
-        stitchingInfo.selectionSetsByType,
+        {},
         stitchingInfo.selectionSetsByField,
         stitchingInfo.dynamicSelectionSetsByField
       ),

packages/delegate/src/results/memoize.ts → packages/delegate/src/memoize.ts

@@ -1,65 +1,118 @@
 import { GraphQLResolveInfo, FieldNode } from 'graphql';
 
-export function memoizeInfoAnd2Objectsand1Primitive<
+export function memoizeInfoAnd2Objects<
   T1 extends GraphQLResolveInfo,
   T2 extends Record<string, any>,
   T3 extends Record<string, any>,
-  T4 extends string,
+  R extends any
+>(fn: (A1: T1, A2: T2, A3: T3) => R): (A1: T1, A2: T2, A3: T3) => R {
+  let cache1: WeakMap<ReadonlyArray<FieldNode>, WeakMap<T2, WeakMap<T3, R>>>;
+
+  function memoized(a1: T1, a2: T2, a3: T3) {
+    if (!cache1) {
+      cache1 = new WeakMap();
+      const cache2: WeakMap<T2, WeakMap<T3, R>> = new WeakMap();
+      cache1.set(a1.fieldNodes, cache2);
+      const cache3: WeakMap<T3, R> = new WeakMap();
+      cache2.set(a2, cache3);
+      const newValue = fn(a1, a2, a3);
+      cache3.set(a3, newValue);
+      return newValue;
+    }
+
+    let cache2 = cache1.get(a1.fieldNodes);
+    if (!cache2) {
+      cache2 = new WeakMap();
+      cache1.set(a1.fieldNodes, cache2);
+      const cache3: WeakMap<T3, R> = new WeakMap();
+      cache2.set(a2, cache3);
+      const newValue = fn(a1, a2, a3);
+      cache3.set(a3, newValue);
+      return newValue;
+    }
+
+    let cache3 = cache2.get(a2);
+    if (!cache3) {
+      cache3 = new WeakMap();
+      cache2.set(a2, cache3);
+      const newValue = fn(a1, a2, a3);
+      cache3.set(a3, newValue);
+      return newValue;
+    }
+
+    const cachedValue = cache3.get(a3);
+    if (cachedValue === undefined) {
+      const newValue = fn(a1, a2, a3);
+      cache3.set(a3, newValue);
+      return newValue;
+    }
+
+    return cachedValue;
+  }
+
+  return memoized;
+}
+
+export function memoize4<
+  T1 extends Record<string, any>,
+  T2 extends Record<string, any>,
+  T3 extends Record<string, any>,
+  T4 extends Record<string, any>,
   R extends any
 >(fn: (A1: T1, A2: T2, A3: T3, A4: T4) => R): (A1: T1, A2: T2, A3: T3, A4: T4) => R {
-  let cache1: WeakMap<ReadonlyArray<FieldNode>, WeakMap<T2, WeakMap<T3, Record<string, R>>>>;
+  let cache1: WeakMap<T1, WeakMap<T2, WeakMap<T3, WeakMap<T4, R>>>>;
 
   function memoized(a1: T1, a2: T2, a3: T3, a4: T4) {
     if (!cache1) {
       cache1 = new WeakMap();
-      const cache2: WeakMap<T2, WeakMap<T3, Record<T4, R>>> = new WeakMap();
-      cache1.set(a1.fieldNodes, cache2);
-      const cache3: WeakMap<T3, Record<T4, R>> = new WeakMap();
+      const cache2: WeakMap<T2, WeakMap<T3, WeakMap<T4, R>>> = new WeakMap();
+      cache1.set(a1, cache2);
+      const cache3: WeakMap<T3, WeakMap<T4, R>> = new WeakMap();
       cache2.set(a2, cache3);
-      const cache4 = Object.create(null);
+      const cache4: WeakMap<T4, R> = new WeakMap();
       cache3.set(a3, cache4);
       const newValue = fn(a1, a2, a3, a4);
-      cache4[a4] = newValue;
+      cache4.set(a4, newValue);
       return newValue;
     }
 
-    let cache2 = cache1.get(a1.fieldNodes);
+    let cache2 = cache1.get(a1);
     if (!cache2) {
       cache2 = new WeakMap();
-      cache1.set(a1.fieldNodes, cache2);
-      const cache3: WeakMap<T3, Record<T4, R>> = new WeakMap();
+      cache1.set(a1, cache2);
+      const cache3: WeakMap<T3, WeakMap<T4, R>> = new WeakMap();
       cache2.set(a2, cache3);
-      const cache4 = Object.create(null);
+      const cache4: WeakMap<T4, R> = new WeakMap();
       cache3.set(a3, cache4);
       const newValue = fn(a1, a2, a3, a4);
-      cache4[a4] = newValue;
+      cache4.set(a4, newValue);
       return newValue;
     }
 
     let cache3 = cache2.get(a2);
     if (!cache3) {
       cache3 = new WeakMap();
       cache2.set(a2, cache3);
-      const cache4 = Object.create(null);
+      const cache4: WeakMap<T4, R> = new WeakMap();
       cache3.set(a3, cache4);
       const newValue = fn(a1, a2, a3, a4);
-      cache4[a4] = newValue;
+      cache4.set(a4, newValue);
       return newValue;
     }
 
-    let cache4 = cache3.get(a3);
+    const cache4 = cache3.get(a3);
     if (!cache4) {
-      cache4 = Object.create(null);
+      const cache4: WeakMap<T4, R> = new WeakMap();
       cache3.set(a3, cache4);
       const newValue = fn(a1, a2, a3, a4);
-      cache4[a4] = newValue;
+      cache4.set(a4, newValue);
       return newValue;
     }
 
-    const cachedValue = cache4[a4];
+    const cachedValue = cache4.get(a4);
     if (cachedValue === undefined) {
       const newValue = fn(a1, a2, a3, a4);
-      cache4[a4] = newValue;
+      cache4.set(a4, newValue);
       return newValue;
     }
 

packages/delegate/src/results/getFieldsNotInSubschema.ts

@@ -4,7 +4,7 @@ import { collectFields, GraphQLExecutionContext } from '@graphql-tools/utils';
 import { isSubschemaConfig } from '../Subschema';
 import { MergedTypeInfo, SubschemaConfig, StitchingInfo } from '../types';
 
-import { memoizeInfoAnd2Objectsand1Primitive } from './memoize';
+import { memoizeInfoAnd2Objects } from '../memoize';
 
 function collectSubFields(info: GraphQLResolveInfo, typeName: string): Record<string, Array<FieldNode>> {
   let subFieldNodes: Record<string, Array<FieldNode>> = Object.create(null);
@@ -28,21 +28,10 @@ function collectSubFields(info: GraphQLResolveInfo, typeName: string): Record<st
   });
 
   const stitchingInfo = info.schema.extensions.stitchingInfo as StitchingInfo;
-  const selectionSetsByType = stitchingInfo.selectionSetsByType;
   const selectionSetsByField = stitchingInfo.selectionSetsByField;
 
   Object.keys(subFieldNodes).forEach(responseName => {
     const fieldName = subFieldNodes[responseName][0].name.value;
-    const typeSelectionSet = selectionSetsByType[typeName];
-    if (typeSelectionSet != null) {
-      subFieldNodes = collectFields(
-        partialExecutionContext,
-        type,
-        typeSelectionSet,
-        subFieldNodes,
-        visitedFragmentNames
-      );
-    }
     const fieldSelectionSet = selectionSetsByField?.[typeName]?.[fieldName];
     if (fieldSelectionSet != null) {
       subFieldNodes = collectFields(
@@ -58,13 +47,13 @@ function collectSubFields(info: GraphQLResolveInfo, typeName: string): Record<st
   return subFieldNodes;
 }
 
-export const getFieldsNotInSubschema = memoizeInfoAnd2Objectsand1Primitive(function (
+export const getFieldsNotInSubschema = memoizeInfoAnd2Objects(function (
   info: GraphQLResolveInfo,
   subschema: GraphQLSchema | SubschemaConfig,
-  mergedTypeInfo: MergedTypeInfo,
-  typeName: string
+  mergedTypeInfo: MergedTypeInfo
 ): Array<FieldNode> {
   const typeMap = isSubschemaConfig(subschema) ? mergedTypeInfo.typeMaps.get(subschema) : subschema.getTypeMap();
+  const typeName = mergedTypeInfo.typeName;
   const fields = (typeMap[typeName] as GraphQLObjectType).getFields();
 
   const subFieldNodes = collectSubFields(info, typeName);

packages/delegate/src/results/handleObject.ts

@@ -43,7 +43,7 @@ export function handleObject(
     return object;
   }
 
-  const fieldNodes = getFieldsNotInSubschema(info, subschema, mergedTypeInfo, typeName);
+  const fieldNodes = getFieldsNotInSubschema(info, subschema, mergedTypeInfo);
 
   return mergeFields(
     mergedTypeInfo,

packages/delegate/src/results/mergeFields.ts

@@ -1,30 +1,49 @@
-import { FieldNode, SelectionNode, Kind, GraphQLResolveInfo, SelectionSetNode } from 'graphql';
+import { FieldNode, SelectionNode, Kind, GraphQLResolveInfo, SelectionSetNode, GraphQLObjectType } from 'graphql';
+
 import isPromise from 'is-promise';
 
+import { typesContainSelectionSet } from '@graphql-tools/utils';
+
 import { MergedTypeInfo, SubschemaConfig } from '../types';
+import { memoize4, memoize3, memoize2 } from '../memoize';
 
-import { memoize3, memoize2 } from './memoize';
 import { mergeProxiedResults } from './mergeProxiedResults';
 
-const sortSubschemasByProxiability = memoize3(function (
+const sortSubschemasByProxiability = memoize4(function (
   mergedTypeInfo: MergedTypeInfo,
   sourceSubschemaOrSourceSubschemas: SubschemaConfig | Array<SubschemaConfig>,
-  targetSubschemas: Array<SubschemaConfig>
+  targetSubschemas: Array<SubschemaConfig>,
+  fieldNodes: Array<FieldNode>
 ): {
   proxiableSubschemas: Array<SubschemaConfig>;
   nonProxiableSubschemas: Array<SubschemaConfig>;
 } {
   // 1.  calculate if possible to delegate to given subschema
-  //    TODO: change logic so that required selection set can be spread across multiple subschemas?
 
   const proxiableSubschemas: Array<SubschemaConfig> = [];
   const nonProxiableSubschemas: Array<SubschemaConfig> = [];
 
   const sourceSubschemas = Array.isArray(sourceSubschemaOrSourceSubschemas)
     ? sourceSubschemaOrSourceSubschemas
     : [sourceSubschemaOrSourceSubschemas];
+
+  const typeName = mergedTypeInfo.typeName;
+  const types = sourceSubschemas.map(sourceSubschema => sourceSubschema.schema.getType(typeName) as GraphQLObjectType);
+
   targetSubschemas.forEach(t => {
-    if (sourceSubschemas.some(s => mergedTypeInfo.containsSelectionSet.get(s).get(t))) {
+    const selectionSet = mergedTypeInfo.selectionSets.get(t);
+    const fieldSelectionSets = mergedTypeInfo.fieldSelectionSets.get(t);
+    if (!typesContainSelectionSet(types, selectionSet)) {
+      nonProxiableSubschemas.push(t);
+    } else if (fieldSelectionSets == null) {
+      proxiableSubschemas.push(t);
+    } else if (
+      fieldNodes.every(fieldNode => {
+        const fieldName = fieldNode.name.value;
+        const fieldSelectionSet = fieldSelectionSets[fieldName];
+        return fieldSelectionSet == null || typesContainSelectionSet(types, fieldSelectionSet);
+      })
+    ) {
       proxiableSubschemas.push(t);
     } else {
       nonProxiableSubschemas.push(t);
@@ -140,7 +159,8 @@ export function mergeFields(
   const { proxiableSubschemas, nonProxiableSubschemas } = sortSubschemasByProxiability(
     mergedTypeInfo,
     sourceSubschemaOrSourceSubschemas,
-    targetSubschemas
+    targetSubschemas,
+    fieldNodes
   );
 
   const { delegationMap, unproxiableFieldNodes } = buildDelegationPlan(mergedTypeInfo, fieldNodes, proxiableSubschemas);