add defer support for inline fragments and fragment spreads

Author: kmoore-intuit

src/main/java/com/intuit/graphql/orchestrator/authorization/DownstreamQueryRedactor.java

@@ -8,10 +8,11 @@
 import graphql.language.Node;
 import graphql.schema.DataFetchingEnvironment;
 import graphql.schema.GraphQLType;
-import java.util.Map;
 import lombok.Builder;
 import lombok.NonNull;
 
+import java.util.Map;
+
 @Builder
 public class DownstreamQueryRedactor {
 
@@ -31,7 +32,9 @@ public DownstreamQueryRedactorResult redact() {
     return new DownstreamQueryRedactorResult(
         transformedRoot,
         downstreamQueryModifier.getDeclineFieldErrors(),
-        downstreamQueryModifier.redactedQueryHasEmptySelectionSet());
+        downstreamQueryModifier.redactedQueryHasEmptySelectionSet(),
+        downstreamQueryModifier.getFragmentSpreadsRemoved()
+      );
   }
 
   private AuthDownstreamQueryModifier createQueryModifier() {

src/main/java/com/intuit/graphql/orchestrator/authorization/DownstreamQueryRedactorResult.java

@@ -2,11 +2,12 @@
 
 import graphql.GraphqlErrorException;
 import graphql.language.Node;
-import java.util.List;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 import lombok.NonNull;
 
+import java.util.List;
+
 @Getter
 @AllArgsConstructor
 public class DownstreamQueryRedactorResult {
@@ -17,4 +18,7 @@ public class DownstreamQueryRedactorResult {
   private List<GraphqlErrorException> errors;
 
   boolean hasEmptySelectionSet;
+
+  @NonNull
+  private List<String> fragmentSpreadsRemoved;
 }

src/main/java/com/intuit/graphql/orchestrator/batch/AuthDownstreamQueryModifier.java

@@ -11,12 +11,14 @@
 import com.intuit.graphql.orchestrator.schema.transform.FieldResolverContext;
 import com.intuit.graphql.orchestrator.utils.SelectionCollector;
 import graphql.GraphQLContext;
+import graphql.GraphQLException;
 import graphql.GraphqlErrorException;
 import graphql.language.Argument;
 import graphql.language.BooleanValue;
 import graphql.language.Directive;
 import graphql.language.Field;
 import graphql.language.FragmentDefinition;
+import graphql.language.FragmentSpread;
 import graphql.language.InlineFragment;
 import graphql.language.Node;
 import graphql.language.NodeVisitorStub;
@@ -74,6 +76,7 @@ public class AuthDownstreamQueryModifier extends NodeVisitorStub {
   private static final ArgumentValueResolver ARGUMENT_VALUE_RESOLVER = new ArgumentValueResolver(); // thread-safe
   private final List<SelectionSetMetadata> processedSelectionSetMetadata = new ArrayList<>();
   private final List<GraphqlErrorException> declinedFieldsErrors = new ArrayList<>();
+  private final List<String> fragmentSpreadsRemoved = new ArrayList<>();
 
   private boolean hasEmptySelectionSet;
 
@@ -105,20 +108,9 @@ public TraversalControl visitField(Field node, TraverserContext<Node> context) {
         return deleteNode(context);
       }
 
-      if(!node.getDirectives(DEFER_DIRECTIVE_NAME).isEmpty()) {
-        Argument deferArg = node.getDirectives(DEFER_DIRECTIVE_NAME).get(0).getArgument(DEFER_IF_ARG);
-        if(graphQLContext.getOrDefault(USE_DEFER, false) && (deferArg == null || ((BooleanValue)deferArg.getValue()).isValue())) {
-          decreaseParentSelectionSetCount(context.getParentContext());
-          return deleteNode(context);
-        } else {
-          //remove directive from query since directive is not built in and will fail downstream if added
-          List<Directive> directives = node.getDirectives()
-                  .stream()
-                  .filter(directive -> !DEFER_DIRECTIVE_NAME.equals(directive.getName()))
-                  .collect(Collectors.toList());
-
-          return changeNode(context, node.transform(builder -> builder.directives(directives)));
-        }
+      List<Directive> directives = node.getDirectives();
+      if(containsDeferDirective(directives)) {
+        return pruneDeferInfo(node, context, directives);
       }
 
       if(!serviceMetadata.getRenamedMetadata().getOriginalFieldNamesByRenamedName().isEmpty()) {
@@ -232,6 +224,24 @@ public TraversalControl visitFragmentDefinition(
   public TraversalControl visitInlineFragment(InlineFragment node, TraverserContext<Node> context) {
     String typeName = node.getTypeCondition().getName();
     context.setVar(GraphQLType.class, this.graphQLSchema.getType(typeName));
+
+    List<Directive> directives = node.getDirectives();
+    if(containsDeferDirective(directives)) {
+      return pruneDeferInfo(node, context, directives);
+    }
+
+    return TraversalControl.CONTINUE;
+  }
+
+  @Override
+  public TraversalControl visitFragmentSpread(FragmentSpread node, TraverserContext<Node> context) {
+    context.setVar(GraphQLType.class, this.graphQLSchema.getType(node.getName()));
+
+    List<Directive> directives = node.getDirectives();
+    if(containsDeferDirective(directives)) {
+      return pruneDeferInfo(node, context, directives);
+    }
+
     return TraversalControl.CONTINUE;
   }
 
@@ -319,10 +329,52 @@ private GraphQLType getParentType(TraverserContext<Node> context) {
     return GraphQLTypeUtil.unwrapAll(parentType);
   }
 
+  private boolean containsDeferDirective(List<Directive> directives) {
+    return directives != null && directives.stream()
+            .anyMatch(directive -> DEFER_DIRECTIVE_NAME.equals(directive.getName()));
+  }
+
+  private TraversalControl pruneDeferInfo(Node node, TraverserContext<Node> context, List<Directive> nodeDirectives) {
+    Directive deferDirective = nodeDirectives
+            .stream()
+            .filter(directive -> DEFER_DIRECTIVE_NAME.equals(directive.getName()))
+            .findFirst()
+            .get();
+
+    Argument deferArg = deferDirective.getArgument(DEFER_IF_ARG);
+    if(graphQLContext.getOrDefault(USE_DEFER, false) && (deferArg == null || ((BooleanValue)deferArg.getValue()).isValue())) {
+      decreaseParentSelectionSetCount(context.getParentContext());
+      if(node instanceof FragmentSpread) {
+        this.fragmentSpreadsRemoved.add(((FragmentSpread)node).getName());
+      }
+
+      return deleteNode(context);
+    } else {
+      final List<Directive> directives = nodeDirectives
+              .stream()
+              .filter(directive -> !DEFER_DIRECTIVE_NAME.equals(directive.getName()))
+              .collect(Collectors.toList());
+      //remove directive from query since directive is not built in and will fail downstream if added
+      if(node instanceof Field) {
+        return changeNode(context, ((Field)node).transform(builder -> builder.directives(directives)));
+      } else if(node instanceof InlineFragment) {
+        return changeNode(context, ((InlineFragment)node).transform(builder -> builder.directives(directives)));
+      } else if(node instanceof FragmentSpread) {
+        return changeNode(context, ((FragmentSpread)node).transform(builder -> builder.directives(directives)));
+      } else {
+        throw new GraphQLException("Not Supported Defer Location.");
+      }
+    }
+  }
+
   public List<GraphqlErrorException> getDeclineFieldErrors() {
     return declinedFieldsErrors;
   }
 
+  public List<String> getFragmentSpreadsRemoved() {
+    return this.fragmentSpreadsRemoved;
+  }
+
   public List<SelectionSetMetadata> getEmptySelectionSets() {
     return this.processedSelectionSetMetadata.stream()
         .filter(selectionSetMetadata -> selectionSetMetadata.getRemainingSelectionsCount() == 0)

src/main/java/com/intuit/graphql/orchestrator/batch/GraphQLServiceBatchLoader.java

@@ -42,6 +42,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -117,15 +118,17 @@ private CompletableFuture<List<DataFetcherResult<Object>>> load(List<DataFetchin
 
     Map<String, FragmentDefinition> mergedFragmentDefinitions = new HashMap<>();
 
+    //Possibly add a complex object to hold these two
     MultiValuedMap<String, GraphqlErrorException> queryRedactErrorsByKey = new ArrayListValuedHashMap<>();
+    Set<String> removedFragments = new HashSet<>();
 
     for (final DataFetchingEnvironment key : keys) {
       MergedFieldModifierResult result = new MergedFieldModifier(key).getFilteredRootField();
       MergedField filteredRootField = result.getMergedField();
       if (filteredRootField != null) {
         filteredRootField.getFields().stream()
           .map(field -> removeFieldsWithExternalTypes(field,
-                operationObjectType, key, authData, fieldAuthorization, queryRedactErrorsByKey, context))
+                operationObjectType, key, authData, fieldAuthorization, queryRedactErrorsByKey, removedFragments, context))
           .filter(Objects::nonNull) // denied access or has an empty selectionSet
           .forEach(selectionSetBuilder::selection);
       }
@@ -195,6 +198,7 @@ private CompletableFuture<List<DataFetcherResult<Object>>> load(List<DataFetchin
         .collect(Collectors.toList());
 
     List<Definition> fragmentsAsDefinitions = mergedFragmentDefinitions.values().stream()
+        .filter(fragment -> !removedFragments.contains(fragment.getName()))
         .map(GraphQLObjects::<Definition>cast).collect(Collectors.toList());
 
     OperationDefinition query = OperationDefinition.newOperationDefinition()
@@ -229,6 +233,7 @@ private CompletableFuture<List<DataFetcherResult<Object>>> load(List<DataFetchin
       query = queryOperationModifier.modifyQuery(graphQLSchema, query, mergedFragmentDefinitions, filteredVariables);
     }
 
+
     return execute(context, query, filteredVariables, fragmentsAsDefinitions)
         .thenApply(queryResponseModifier::modify)
         .thenApply(result -> batchResultTransformer.toBatchResult(result, keys))
@@ -370,7 +375,7 @@ private FragmentDefinition removeDomainTypeFromFragment(final FragmentDefinition
 
   private Field removeFieldsWithExternalTypes(Field origField, GraphQLObjectType operationObjectType,
       DataFetchingEnvironment dfe, Object authData, FieldAuthorization fieldAuthorization,
-      final MultiValuedMap<String, GraphqlErrorException> queryRedactErrorsByKey, GraphQLContext context) {
+      final MultiValuedMap<String, GraphqlErrorException> queryRedactErrorsByKey, Set<String> removedFragments, GraphQLContext context) {
 
     if (serviceMetadata.shouldModifyDownStreamQuery() ||
         !(fieldAuthorization instanceof DefaultFieldAuthorization) ||
@@ -391,6 +396,10 @@ private Field removeFieldsWithExternalTypes(Field origField, GraphQLObjectType o
         String keyPath = dfe.getExecutionStepInfo().getPath().toString();
         queryRedactErrorsByKey.putAll(keyPath, redactResult.getErrors());
       }
+      if (CollectionUtils.isNotEmpty(redactResult.getFragmentSpreadsRemoved())) {
+        removedFragments.addAll(redactResult.getFragmentSpreadsRemoved());
+      }
+
       if (redactResult.isHasEmptySelectionSet()) {
         return null; // if null, not added to selection
       }

src/main/java/com/intuit/graphql/orchestrator/deferDirective/MultipartQueryRequest.java

@@ -0,0 +1,16 @@
+package com.intuit.graphql.orchestrator.deferDirective;
+
+import graphql.language.OperationDefinition;
+import lombok.Builder;
+import lombok.Getter;
+import lombok.Singular;
+
+import java.util.Set;
+
+@Getter
+@Builder
+public class MultipartQueryRequest {
+    private OperationDefinition multipartOperationDef;
+    @Singular
+    private Set<String> fragmentSpreadNames;
+}

src/main/java/com/intuit/graphql/orchestrator/utils/DirectivesUtil.java

@@ -87,12 +87,18 @@ private static GraphQLDirective createDeferDirective() {
             .defaultValueLiteral(BooleanValue.newBooleanValue(true).build())
             .build();
 
+    GraphQLArgument labelArgument = GraphQLArgument.newArgument()
+            .name(DEFER_IF_ARG)
+            .type(Scalars.GraphQLString)
+            .build();
+
     return GraphQLDirective.newDirective()
             .name(DEFER_DIRECTIVE_NAME)
             .validLocation(DirectiveLocation.FIELD)
             .validLocation(DirectiveLocation.INLINE_FRAGMENT)
             .validLocation(DirectiveLocation.FRAGMENT_SPREAD)
             .argument(ifArgument)
+            .argument(labelArgument)
             .build();
   }