fixes selectionSet hints with WrapFields transform (#1757)

Visitors visiting the proxying ast cannot rely on the operation being valid according to the source schema in all cases.

But, the proxying document's root field nodes represent values with the expected return type, and so the document can be traversed according to the source schema from that point and below.

See #1725

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);
+    });
   });
 });