fix combination of type merging and gateway type extensions

packages/delegate/src/mergeFields.ts

@@ -5,12 +5,12 @@ import { MergedTypeInfo, SubschemaConfig } from './types';
 
 function buildDelegationPlan(
   mergedTypeInfo: MergedTypeInfo,
-  originalSelections: Array<FieldNode>,
+  fieldNodes: Array<FieldNode>,
   sourceSubschemas: Array<SubschemaConfig>,
   targetSubschemas: Array<SubschemaConfig>
 ): {
   delegationMap: Map<SubschemaConfig, Array<SelectionNode>>;
-  unproxiableSelections: Array<FieldNode>;
+  unproxiableFieldNodes: Array<FieldNode>;
   proxiableSubschemas: Array<SubschemaConfig>;
   nonProxiableSubschemas: Array<SubschemaConfig>;
 } {
@@ -34,53 +34,62 @@ function buildDelegationPlan(
   });
 
   const { uniqueFields, nonUniqueFields } = mergedTypeInfo;
-  const unproxiableSelections: Array<FieldNode> = [];
+  const unproxiableFieldNodes: Array<FieldNode> = [];
 
   // 2. for each selection:
 
   const delegationMap: Map<SubschemaConfig, Array<SelectionNode>> = new Map();
-  originalSelections.forEach(selection => {
-    if (selection.name.value === '__typename') {
+  fieldNodes.forEach(fieldNode => {
+    if (fieldNode.name.value === '__typename') {
       return;
     }
 
     // 2a. use uniqueFields map to assign fields to subschema if one of possible subschemas
 
-    const uniqueSubschema: SubschemaConfig = uniqueFields[selection.name.value];
+    const uniqueSubschema: SubschemaConfig = uniqueFields[fieldNode.name.value];
     if (uniqueSubschema != null) {
-      if (proxiableSubschemas.includes(uniqueSubschema)) {
-        const existingSubschema = delegationMap.get(uniqueSubschema);
-        if (existingSubschema != null) {
-          existingSubschema.push(selection);
-        } else {
-          delegationMap.set(uniqueSubschema, [selection]);
-        }
-      } else {
-        unproxiableSelections.push(selection);
+      if (!proxiableSubschemas.includes(uniqueSubschema)) {
+        unproxiableFieldNodes.push(fieldNode);
+        return;
       }
-    } else {
-      // 2b. use nonUniqueFields to assign to a possible subschema,
-      //     preferring one of the subschemas already targets of delegation
-
-      let nonUniqueSubschemas: Array<SubschemaConfig> = nonUniqueFields[selection.name.value];
-      nonUniqueSubschemas = nonUniqueSubschemas.filter(s => proxiableSubschemas.includes(s));
-      if (nonUniqueSubschemas != null) {
-        const subschemas: Array<SubschemaConfig> = Array.from(delegationMap.keys());
-        const existingSubschema = nonUniqueSubschemas.find(s => subschemas.includes(s));
-        if (existingSubschema != null) {
-          delegationMap.get(existingSubschema).push(selection);
-        } else {
-          delegationMap.set(nonUniqueSubschemas[0], [selection]);
-        }
+
+      const existingSubschema = delegationMap.get(uniqueSubschema);
+      if (existingSubschema != null) {
+        existingSubschema.push(fieldNode);
       } else {
-        unproxiableSelections.push(selection);
+        delegationMap.set(uniqueSubschema, [fieldNode]);
       }
+
+      return;
+    }
+
+    // 2b. use nonUniqueFields to assign to a possible subschema,
+    //     preferring one of the subschemas already targets of delegation
+
+    let nonUniqueSubschemas: Array<SubschemaConfig> = nonUniqueFields[fieldNode.name.value];
+    if (nonUniqueSubschemas == null) {
+      unproxiableFieldNodes.push(fieldNode);
+      return;
+    }
+
+    nonUniqueSubschemas = nonUniqueSubschemas.filter(s => proxiableSubschemas.includes(s));
+    if (nonUniqueSubschemas == null) {
+      unproxiableFieldNodes.push(fieldNode);
+      return;
+    }
+
+    const subschemas: Array<SubschemaConfig> = Array.from(delegationMap.keys());
+    const existingSubschema = nonUniqueSubschemas.find(s => subschemas.includes(s));
+    if (existingSubschema != null) {
+      delegationMap.get(existingSubschema).push(fieldNode);
+    } else {
+      delegationMap.set(nonUniqueSubschemas[0], [fieldNode]);
     }
   });
 
   return {
     delegationMap,
-    unproxiableSelections,
+    unproxiableFieldNodes,
     proxiableSubschemas,
     nonProxiableSubschemas,
   };
@@ -90,19 +99,19 @@ export function mergeFields(
   mergedTypeInfo: MergedTypeInfo,
   typeName: string,
   object: any,
-  originalSelections: Array<FieldNode>,
+  fieldNodes: Array<FieldNode>,
   sourceSubschemas: Array<SubschemaConfig>,
   targetSubschemas: Array<SubschemaConfig>,
   context: Record<string, any>,
   info: GraphQLResolveInfo
 ): any {
-  if (!originalSelections.length) {
+  if (!fieldNodes.length) {
     return object;
   }
 
-  const { delegationMap, unproxiableSelections, proxiableSubschemas, nonProxiableSubschemas } = buildDelegationPlan(
+  const { delegationMap, unproxiableFieldNodes, proxiableSubschemas, nonProxiableSubschemas } = buildDelegationPlan(
     mergedTypeInfo,
-    originalSelections,
+    fieldNodes,
     sourceSubschemas,
     targetSubschemas
   );
@@ -134,7 +143,7 @@ export function mergeFields(
           mergedTypeInfo,
           typeName,
           mergeProxiedResults(object, ...results),
-          unproxiableSelections,
+          unproxiableFieldNodes,
           sourceSubschemas.concat(proxiableSubschemas),
           nonProxiableSubschemas,
           context,
@@ -145,7 +154,7 @@ export function mergeFields(
         mergedTypeInfo,
         typeName,
         mergeProxiedResults(object, ...maybePromises),
-        unproxiableSelections,
+        unproxiableFieldNodes,
         sourceSubschemas.concat(proxiableSubschemas),
         nonProxiableSubschemas,
         context,

packages/delegate/src/results/handleObject.ts

@@ -11,7 +11,7 @@ import {
 import { collectFields, GraphQLExecutionContext, setErrors, slicedError } from '@graphql-tools/utils';
 import { setObjectSubschema, isSubschemaConfig } from '../Subschema';
 import { mergeFields } from '../mergeFields';
-import { MergedTypeInfo, SubschemaConfig } from '../types';
+import { MergedTypeInfo, SubschemaConfig, StitchingInfo } from '../types';
 
 export function handleObject(
   type: GraphQLCompositeType,
@@ -52,57 +52,76 @@ export function handleObject(
     return object;
   }
 
-  const subFields = collectSubFields(info, object.__typename);
-
-  const selections = getFieldsNotInSubschema(subFields, subschema, mergedTypeInfo, object.__typename);
+  const fieldNodes = getFieldsNotInSubschema(info, subschema, mergedTypeInfo, object.__typename);
 
   return mergeFields(
     mergedTypeInfo,
     typeName,
     object,
-    selections,
+    fieldNodes,
     [subschema as SubschemaConfig],
     targetSubschemas,
     context,
     info
   );
 }
 
-function collectSubFields(info: GraphQLResolveInfo, typeName: string) {
+function collectSubFields(info: GraphQLResolveInfo, typeName: string): Record<string, Array<FieldNode>> {
   let subFieldNodes: Record<string, Array<FieldNode>> = Object.create(null);
   const visitedFragmentNames = Object.create(null);
+
+  const type = info.schema.getType(typeName) as GraphQLObjectType;
+  const partialExecutionContext = ({
+    schema: info.schema,
+    variableValues: info.variableValues,
+    fragments: info.fragments,
+  } as unknown) as GraphQLExecutionContext;
+
   info.fieldNodes.forEach(fieldNode => {
     subFieldNodes = collectFields(
-      ({
-        schema: info.schema,
-        variableValues: info.variableValues,
-        fragments: info.fragments,
-      } as unknown) as GraphQLExecutionContext,
-      info.schema.getType(typeName) as GraphQLObjectType,
+      partialExecutionContext,
+      type,
       fieldNode.selectionSet,
       subFieldNodes,
       visitedFragmentNames
     );
   });
+
+  const selectionSetsByField = (info.schema.extensions.stitchingInfo as StitchingInfo).selectionSetsByField;
+
+  Object.keys(subFieldNodes).forEach(responseName => {
+    const fieldName = subFieldNodes[responseName][0].name.value;
+    if (selectionSetsByField[typeName] && selectionSetsByField[typeName][fieldName]) {
+      subFieldNodes = collectFields(
+        partialExecutionContext,
+        type,
+        selectionSetsByField[typeName][fieldName],
+        subFieldNodes,
+        visitedFragmentNames
+      );
+    }
+  });
+
   return subFieldNodes;
 }
 
 function getFieldsNotInSubschema(
-  subFieldNodes: Record<string, Array<FieldNode>>,
+  info: GraphQLResolveInfo,
   subschema: GraphQLSchema | SubschemaConfig,
   mergedTypeInfo: MergedTypeInfo,
   typeName: string
 ): Array<FieldNode> {
   const typeMap = isSubschemaConfig(subschema) ? mergedTypeInfo.typeMaps.get(subschema) : subschema.getTypeMap();
   const fields = (typeMap[typeName] as GraphQLObjectType).getFields();
 
-  const fieldsNotInSchema: Array<FieldNode> = [];
+  const subFieldNodes = collectSubFields(info, typeName);
+
+  let fieldsNotInSchema: Array<FieldNode> = [];
   Object.keys(subFieldNodes).forEach(responseName => {
-    subFieldNodes[responseName].forEach(subFieldNode => {
-      if (!(subFieldNode.name.value in fields)) {
-        fieldsNotInSchema.push(subFieldNode);
-      }
-    });
+    const fieldName = subFieldNodes[responseName][0].name.value;
+    if (!(fieldName in fields)) {
+      fieldsNotInSchema = fieldsNotInSchema.concat(subFieldNodes[responseName]);
+    }
   });
 
   return fieldsNotInSchema;

packages/stitch/tests/typeMerging.test.ts

@@ -7,6 +7,8 @@ import { makeExecutableSchema } from '@graphql-tools/schema';
 
 import { addMocksToSchema } from '@graphql-tools/mock';
 
+import { delegateToSchema } from '@graphql-tools/delegate';
+
 import { stitchSchemas } from '../src/stitchSchemas';
 
 let chirpSchema = makeExecutableSchema({
@@ -103,3 +105,133 @@ describe('merging using type merging', () => {
     expect(result.data.userById.chirps[1].author.email).not.toBe(null);
   });
 });
+
+describe('merge types and extend', () => {
+  test('should work', async () => {
+    const resultSchema = makeExecutableSchema({
+      typeDefs: `
+        type Query {
+          resultById(id: ID!): String
+        }
+      `,
+      resolvers: {
+        Query: {
+          resultById: () => 'ok',
+        },
+      },
+    });
+
+    const containerSchemaA = makeExecutableSchema({
+      typeDefs: `
+          type Container {
+            id: ID!
+            resultId: ID!
+          }
+
+          type Query {
+            containerById(id: ID!): Container
+          }
+      `,
+      resolvers: {
+        Query: {
+          containerById: () => ({ id: 'Container', resultId: 'Result' }),
+        },
+      },
+    });
+
+    const containerSchemaB = makeExecutableSchema({
+      typeDefs: `
+        type Container {
+          id: ID!
+        }
+
+        type Query {
+          containerById(id: ID!): Container
+          rootContainer: Container!
+        }
+      `,
+      resolvers: {
+        Query: {
+          containerById: () => ({ id: 'Container' }),
+          rootContainer: () => ({ id: 'Container' }),
+        },
+      },
+    });
+
+    const schema = stitchSchemas({
+      subschemas: [
+        {
+          schema: resultSchema,
+        },
+        {
+          schema: containerSchemaA,
+          merge: {
+            Container: {
+              fieldName: 'containerById',
+              args: ({ id }) => ({ id }),
+              selectionSet: '{ id }',
+            },
+          },
+        },
+        {
+          schema: containerSchemaB,
+          merge: {
+            Container: {
+              fieldName: 'containerById',
+              args: ({ id }) => ({ id }),
+              selectionSet: '{ id }',
+            },
+          },
+        },
+      ],
+      mergeTypes: true,
+      typeDefs: `
+        extend type Container {
+          result: String!
+        }
+      `,
+      resolvers: {
+        Container: {
+          result: {
+            selectionSet: `{ resultId }`,
+            resolve(container, _args, context, info) {
+              return delegateToSchema({
+                schema: resultSchema,
+                operation: 'query',
+                fieldName: 'resultById',
+                args: {
+                  id: container.resultId,
+                },
+                context,
+                info,
+              });
+            },
+          },
+        },
+      },
+    });
+
+    const result = await graphql(
+      schema,
+      `
+        query TestQuery {
+          rootContainer {
+            id
+            result
+          }
+        }
+      `,
+    );
+
+    const expectedResult = {
+      data: {
+        rootContainer: {
+          id: 'Container',
+          result: 'ok',
+        }
+      }
+    }
+
+    expect(result).toEqual(expectedResult);
+  })
+})