fixes selectionSet hints with WrapFields transform

packages/delegate/src/delegationBindings.ts

@@ -2,12 +2,11 @@ import { Transform } from '@graphql-tools/utils';
 
 import { StitchingInfo, DelegationContext } from './types';
 
+import AddSelectionSets from './transforms/AddSelectionSets';
 import ExpandAbstractTypes from './transforms/ExpandAbstractTypes';
 import WrapConcreteTypes from './transforms/WrapConcreteTypes';
 import FilterToSchema from './transforms/FilterToSchema';
 import AddFragmentsByField from './transforms/AddFragmentsByField';
-import AddSelectionSetsByField from './transforms/AddSelectionSetsByField';
-import AddSelectionSetsByType from './transforms/AddSelectionSetsByType';
 import AddTypenameToAbstract from './transforms/AddTypenameToAbstract';
 import CheckResultAndHandleErrors from './transforms/CheckResultAndHandleErrors';
 import AddArgumentsAsVariables from './transforms/AddArgumentsAsVariables';
@@ -39,8 +38,12 @@ export function defaultDelegationBinding(delegationContext: DelegationContext):
 
   if (stitchingInfo != null) {
     delegationTransforms = delegationTransforms.concat([
-      new AddSelectionSetsByField(info.schema, stitchingInfo.selectionSetsByField),
-      new AddSelectionSetsByType(info.schema, stitchingInfo.selectionSetsByType),
+      new AddSelectionSets(
+        info.schema,
+        stitchingInfo.selectionSetsByType,
+        stitchingInfo.selectionSetsByField,
+        returnType
+      ),
       new WrapConcreteTypes(returnType, transformedSchema),
       new ExpandAbstractTypes(info.schema, transformedSchema),
     ]);

packages/delegate/src/transforms/AddSelectionSets.ts

@@ -0,0 +1,62 @@
+import { GraphQLSchema, SelectionSetNode, TypeInfo, GraphQLOutputType, Kind } from 'graphql';
+
+import { Transform, Request } from '@graphql-tools/utils';
+import VisitSelectionSets from './VisitSelectionSets';
+
+export default class AddSelectionSetsByField implements Transform {
+  private readonly transformer: VisitSelectionSets;
+
+  constructor(
+    sourceSchema: GraphQLSchema,
+    selectionSetsByType: Record<string, SelectionSetNode>,
+    selectionSetsByField: Record<string, Record<string, SelectionSetNode>>,
+    initialType: GraphQLOutputType
+  ) {
+    this.transformer = new VisitSelectionSets(sourceSchema, initialType, (node, typeInfo) =>
+      visitSelectionSet(node, typeInfo, selectionSetsByType, selectionSetsByField)
+    );
+  }
+
+  public transformRequest(originalRequest: Request): Request {
+    return this.transformer.transformRequest(originalRequest);
+  }
+}
+
+function visitSelectionSet(
+  node: SelectionSetNode,
+  typeInfo: TypeInfo,
+  selectionSetsByType: Record<string, SelectionSetNode>,
+  selectionSetsByField: Record<string, Record<string, SelectionSetNode>>
+): SelectionSetNode {
+  const parentType = typeInfo.getParentType();
+  if (parentType != null) {
+    const parentTypeName = parentType.name;
+    let selections = node.selections;
+
+    if (parentTypeName in selectionSetsByType) {
+      const selectionSet = selectionSetsByType[parentTypeName];
+      if (selectionSet != null) {
+        selections = selections.concat(selectionSet.selections);
+      }
+    }
+
+    if (parentTypeName in selectionSetsByField) {
+      node.selections.forEach(selection => {
+        if (selection.kind === Kind.FIELD) {
+          const name = selection.name.value;
+          const selectionSet = selectionSetsByField[parentTypeName][name];
+          if (selectionSet != null) {
+            selections = selections.concat(selectionSet.selections);
+          }
+        }
+      });
+    }
+
+    if (selections !== node.selections) {
+      return {
+        ...node,
+        selections,
+      };
+    }
+  }
+}

packages/delegate/src/transforms/VisitSelectionSets.ts

@@ -0,0 +1,137 @@
+import {
+  DocumentNode,
+  GraphQLSchema,
+  Kind,
+  SelectionSetNode,
+  TypeInfo,
+  visit,
+  visitWithTypeInfo,
+  GraphQLOutputType,
+  OperationDefinitionNode,
+  FragmentDefinitionNode,
+  SelectionNode,
+  DefinitionNode,
+} from 'graphql';
+
+import { Transform, Request, collectFields, GraphQLExecutionContext } from '@graphql-tools/utils';
+
+export default class VisitSelectionSets implements Transform {
+  private readonly schema: GraphQLSchema;
+  private readonly initialType: GraphQLOutputType;
+  private readonly visitor: (node: SelectionSetNode, typeInfo: TypeInfo) => SelectionSetNode;
+
+  constructor(
+    schema: GraphQLSchema,
+    initialType: GraphQLOutputType,
+    visitor: (node: SelectionSetNode, typeInfo: TypeInfo) => SelectionSetNode
+  ) {
+    this.schema = schema;
+    this.initialType = initialType;
+    this.visitor = visitor;
+  }
+
+  public transformRequest(originalRequest: Request): Request {
+    const document = visitSelectionSets(originalRequest, this.schema, this.initialType, this.visitor);
+    return {
+      ...originalRequest,
+      document,
+    };
+  }
+}
+
+function visitSelectionSets(
+  request: Request,
+  schema: GraphQLSchema,
+  initialType: GraphQLOutputType,
+  visitor: (node: SelectionSetNode, typeInfo: TypeInfo) => SelectionSetNode
+): DocumentNode {
+  const { document, variables } = request;
+
+  const operations: Array<OperationDefinitionNode> = [];
+  const fragments: Record<string, FragmentDefinitionNode> = Object.create(null);
+  document.definitions.forEach(def => {
+    if (def.kind === Kind.OPERATION_DEFINITION) {
+      operations.push(def);
+    } else if (def.kind === Kind.FRAGMENT_DEFINITION) {
+      fragments[def.name.value] = def;
+    }
+  });
+
+  const partialExecutionContext = {
+    schema,
+    variableValues: variables,
+    fragments,
+  } as GraphQLExecutionContext;
+
+  const typeInfo = new TypeInfo(schema, undefined, initialType);
+  const newDefinitions: Array<DefinitionNode> = operations.map(operation => {
+    const type =
+      operation.operation === 'query'
+        ? schema.getQueryType()
+        : operation.operation === 'mutation'
+        ? schema.getMutationType()
+        : schema.getSubscriptionType();
+
+    const fields = collectFields(
+      partialExecutionContext,
+      type,
+      operation.selectionSet,
+      Object.create(null),
+      Object.create(null)
+    );
+
+    const newSelections: Array<SelectionNode> = [];
+    Object.keys(fields).forEach(responseKey => {
+      const fieldNodes = fields[responseKey];
+      fieldNodes.forEach(fieldNode => {
+        const selectionSet = fieldNode.selectionSet;
+
+        if (selectionSet == null) {
+          newSelections.push(fieldNode);
+          return;
+        }
+
+        const newSelectionSet = visit(
+          selectionSet,
+          visitWithTypeInfo(typeInfo, {
+            [Kind.SELECTION_SET]: node => visitor(node, typeInfo),
+          })
+        );
+
+        if (newSelectionSet === selectionSet) {
+          newSelections.push(fieldNode);
+          return;
+        }
+
+        newSelections.push({
+          ...fieldNode,
+          selectionSet: newSelectionSet,
+        });
+      });
+    });
+
+    return {
+      ...operation,
+      selectionSet: {
+        kind: Kind.SELECTION_SET,
+        selections: newSelections,
+      },
+    };
+  });
+
+  Object.values(fragments).forEach(fragment => {
+    newDefinitions.push(
+      visit(
+        fragment,
+        visitWithTypeInfo(typeInfo, {
+          [Kind.SELECTION_SET]: node => visitor(node, typeInfo),
+        })
+      )
+    );
+  });
+
+  return {
+    ...document,
+    definitions: newDefinitions,
+  };
+}

packages/delegate/src/transforms/index.ts

@@ -1,13 +1,14 @@
 export { default as CheckResultAndHandleErrors } from './CheckResultAndHandleErrors';
 export { checkResultAndHandleErrors } from './CheckResultAndHandleErrors';
 export { default as ExpandAbstractTypes } from './ExpandAbstractTypes';
-export { default as AddSelectionSetsByField } from './AddSelectionSetsByField';
-export { default as AddMergedTypeSelectionSets } from './AddSelectionSetsByType';
+export { default as VisitSelectionSets } from './VisitSelectionSets';
+export { default as AddSelectionSets } from './AddSelectionSets';
 export { default as AddArgumentsAsVariables } from './AddArgumentsAsVariables';
 export { default as FilterToSchema } from './FilterToSchema';
 export { default as AddTypenameToAbstract } from './AddTypenameToAbstract';
 
-// superseded by AddFragmentsByField
+// superseded by VisitSelectionSets and AddSelectionSets
+export { default as AddSelectionSetsByField } from './AddSelectionSetsByField';
+export { default as AddMergedTypeSelectionSets } from './AddSelectionSetsByType';
 export { default as ReplaceFieldWithFragment } from './ReplaceFieldWithFragment';
-// superseded by AddSelectionSetsByField
 export { default as AddFragmentsByField } from './AddFragmentsByField';

packages/stitch/tests/alternateStitchSchemas.test.ts

@@ -36,6 +36,7 @@ import {
 } from '@graphql-tools/delegate';
 
 import { makeExecutableSchema } from '@graphql-tools/schema';
+import { addMocksToSchema } from '@graphql-tools/mock';
 import {
   wrapFieldNode,
   renameFieldNode,
@@ -1437,6 +1438,46 @@ describe('schema transformation with wrapping of object fields', () => {
 
       expect(result).toEqual(expectedResult);
     });
+
+    test('should work with selectionSets', async () => {
+      let subschema = makeExecutableSchema({
+        typeDefs: `
+          type Query {
+            user: User
+          }
+          type User {
+            id: ID
+          }
+        `,
+      });
+
+      subschema = addMocksToSchema({ schema: subschema });
+      const stitchedSchema = stitchSchemas({
+        subschemas: [{
+          schema: subschema,
+          transforms: [
+            new WrapFields('Query', ['wrapped'], [`WrappedQuery`]),
+          ],
+        }],
+        typeDefs: `
+          extend type User {
+            dummy: String
+          }
+        `,
+        resolvers: {
+          User: {
+            dummy: {
+              selectionSet: `{ id }`,
+              resolve: (user: any) => user.id,
+            },
+          },
+        },
+      });
+
+      const query = '{ wrapped { user { dummy } } }';
+      const result = await graphql(stitchedSchema, query);
+      expect(result.data.wrapped.user.dummy).not.toEqual(null);
+    });
   });
 });