Defer query visitor

src/main/java/com/intuit/graphql/orchestrator/GraphQLOrchestrator.java

@@ -1,6 +1,7 @@
 package com.intuit.graphql.orchestrator;
 
 import com.intuit.graphql.orchestrator.deferDirective.DeferDirectiveInstrumentation;
+import com.intuit.graphql.orchestrator.deferDirective.DeferOptions;
 import com.intuit.graphql.orchestrator.schema.RuntimeGraph;
 import com.intuit.graphql.orchestrator.utils.MultiEIGenerator;
 import graphql.ExecutionInput;
@@ -41,6 +42,8 @@
 public class GraphQLOrchestrator {
 
   public static final String DATA_LOADER_REGISTRY_CONTEXT_KEY = DataLoaderRegistry.class.getName() + ".context.key";
+  private static final DeferOptions DEFAULT_DEFER_OPTIONS = DeferOptions.builder().nestedDefersAllowed(false).build();
+  private static final boolean DISABLED_DEFER = false;
 
   private final RuntimeGraph runtimeGraph;
   private final List<Instrumentation> instrumentations;
@@ -80,12 +83,12 @@ private DataLoaderRegistry buildNewDataLoaderRegistry() {
   }
 
   public CompletableFuture<ExecutionResult> execute(ExecutionInput executionInput) {
-    return execute(executionInput, false);
+    return execute(executionInput, DEFAULT_DEFER_OPTIONS, DISABLED_DEFER);
   }
 
-  public CompletableFuture<ExecutionResult> execute(ExecutionInput executionInput, boolean hasDefer) {
+  public CompletableFuture<ExecutionResult> execute(ExecutionInput executionInput, DeferOptions deferOptions, boolean hasDefer) {
     if(hasDefer) {
-      return executeWithDefer(executionInput);
+      return executeWithDefer(executionInput, deferOptions);
     }
       final GraphQL graphQL = constructGraphQL();
 
@@ -101,9 +104,9 @@ public CompletableFuture<ExecutionResult> execute(ExecutionInput executionInput,
       return graphQL.executeAsync(newExecutionInput);
   }
 
-  private CompletableFuture<ExecutionResult> executeWithDefer(ExecutionInput executionInput) {
+  private CompletableFuture<ExecutionResult> executeWithDefer(ExecutionInput executionInput, DeferOptions options) {
     AtomicInteger responses = new AtomicInteger(0);
-    MultiEIGenerator eiGenerator = new MultiEIGenerator(executionInput);
+    MultiEIGenerator eiGenerator = new MultiEIGenerator(executionInput, options, this.getSchema());
 
     Flux<Object> executionResultPublisher = eiGenerator.generateEIs()
             .filter(ei -> !ei.getQuery().equals(""))

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

@@ -3,7 +3,6 @@
 import com.intuit.graphql.orchestrator.batch.AuthDownstreamQueryModifier;
 import com.intuit.graphql.orchestrator.schema.ServiceMetadata;
 import com.intuit.graphql.orchestrator.utils.SelectionCollector;
-import graphql.language.AstTransformer;
 import graphql.language.FragmentDefinition;
 import graphql.language.Node;
 import graphql.schema.DataFetchingEnvironment;
@@ -13,11 +12,10 @@
 
 import java.util.Map;
 
+import static com.intuit.graphql.orchestrator.utils.GraphQLUtil.AST_TRANSFORMER;
+
 @Builder
 public class DownstreamQueryRedactor {
-
-  private static final AstTransformer AST_TRANSFORMER = new AstTransformer();
-
   @NonNull private Node<?> root;
   @NonNull private GraphQLType rootType;
   @NonNull private GraphQLType rootParentType;

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

@@ -0,0 +1,10 @@
+package com.intuit.graphql.orchestrator.deferDirective;
+
+import lombok.Builder;
+import lombok.Getter;
+
+@Builder
+@Getter
+public class DeferOptions {
+    private boolean nestedDefersAllowed;
+}

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

@@ -0,0 +1,53 @@
+package com.intuit.graphql.orchestrator.deferDirective;
+
+import graphql.language.Argument;
+import graphql.language.BooleanValue;
+import graphql.language.Directive;
+import graphql.language.DirectivesContainer;
+import graphql.language.Node;
+import graphql.language.Selection;
+import graphql.language.SelectionSet;
+import lombok.NonNull;
+
+import java.util.List;
+
+import static com.intuit.graphql.orchestrator.utils.DirectivesUtil.DEFER_DIRECTIVE_NAME;
+import static com.intuit.graphql.orchestrator.utils.DirectivesUtil.DEFER_IF_ARG;
+
+public class DeferUtil {
+    /**
+     * Checks if it is necessary to create ei for deferred field.
+     * Currently, selection should be skipped if all the children field are deferred resulting in an empty selection set.
+     * @param selection: node to check if children are all deferred
+     * @return boolean: true if all children are deferred, false otherwise
+     */
+    public static boolean hasNonDeferredSelection(@NonNull  Selection selection) {
+        return ((List<Node>)selection.getChildren())
+                .stream()
+                .filter(SelectionSet.class::isInstance)
+                .map(SelectionSet.class::cast)
+                .findAny()
+                .get()
+                .getSelections()
+                .stream()
+                .anyMatch(child -> !containsEnabledDeferDirective(child));
+    }
+
+    /**
+     * Verifies that Node has defer directive that is not disabled
+     * @param node: node to check if contains defer directive
+     * @return boolean: true if node has an enabled defer, false otherwise
+     */
+    public static boolean containsEnabledDeferDirective(Selection node) {
+        return node instanceof DirectivesContainer &&
+                ((List<Directive>) ((DirectivesContainer) node).getDirectives())
+                .stream()
+                .filter(directive ->  DEFER_DIRECTIVE_NAME.equals(directive.getName()))
+                .findFirst()
+                .map(directive -> {
+                    Argument ifArg = directive.getArgument(DEFER_IF_ARG);
+                    return ifArg == null || ((BooleanValue) ifArg.getValue()).isValue();
+                })
+                .orElse(false);
+    }
+}

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

@@ -1,25 +1,41 @@
 package com.intuit.graphql.orchestrator.utils;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.intuit.graphql.orchestrator.deferDirective.DeferOptions;
+import com.intuit.graphql.orchestrator.visitors.queryVisitors.DeferQueryExtractor;
 import graphql.ExecutionInput;
+import graphql.analysis.QueryTransformer;
+import graphql.language.Document;
+import graphql.language.FragmentDefinition;
+import graphql.language.OperationDefinition;
+import graphql.language.SelectionSet;
+import graphql.schema.GraphQLSchema;
 import lombok.extern.slf4j.Slf4j;
 import reactor.core.publisher.Flux;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import static com.intuit.graphql.orchestrator.utils.GraphQLUtil.parser;
 
 @Slf4j
 public class MultiEIGenerator {
-
-    private List<ExecutionInput> eis = new ArrayList<>();
+    private final List<ExecutionInput> eis = new ArrayList<>();
+    private final DeferOptions deferOptions;
+    private final GraphQLSchema schema;
     private Integer numOfEIs = null;
-    private final static String EMPTY_QUERY ="";
 
     @VisibleForTesting
     private long timeProcessedSplit = 0;
 
-    public MultiEIGenerator(ExecutionInput ei) {
+    public MultiEIGenerator(ExecutionInput ei, DeferOptions deferOptions, GraphQLSchema schema) {
         this.eis.add(ei);
+        this.deferOptions = deferOptions;
+        this.schema = schema;
     }
 
     public Flux<ExecutionInput> generateEIs() {
@@ -39,7 +55,40 @@ public Flux<ExecutionInput> generateEIs() {
                 this.timeProcessedSplit = System.currentTimeMillis();
                 //Adds elements to list of eis that need to be processed
                 try {
-                    this.eis.addAll(MultipartUtil.splitMultipartExecutionInput(emittedEI));
+                    Document rootDocument = parser.parseDocument(emittedEI.getQuery());
+                    Map<String, FragmentDefinition> fragmentDefinitionMap = rootDocument.getDefinitionsOfType(FragmentDefinition.class)
+                            .stream()
+                            .collect(Collectors.toMap(FragmentDefinition::getName , Function.identity()));
+
+                    ExecutionInput finalEmittedEI = emittedEI;
+                    AtomicReference<OperationDefinition> operationDefinitionReference = new AtomicReference<>();
+                    rootDocument.getDefinitionsOfType(OperationDefinition.class)
+                    .stream()
+                    .peek(operationDefinitionReference::set)
+                    .map(OperationDefinition::getSelectionSet)
+                    .map(SelectionSet::getSelections)
+                    .flatMap(List::stream)
+                    .forEach(selection -> {
+                        QueryTransformer transformer = QueryTransformer.newQueryTransformer()
+                                .schema(this.schema)
+                                .root(selection)
+                                .rootParentType(this.schema.getQueryType())
+                                .fragmentsByName(fragmentDefinitionMap)
+                                .variables(finalEmittedEI.getVariables())
+                                .build();
+
+                        DeferQueryExtractor visitor = DeferQueryExtractor.builder()
+                                .deferOptions(deferOptions)
+                                .originalEI(finalEmittedEI)
+                                .rootNode(rootDocument)
+                                .operationDefinition(operationDefinitionReference.get())
+                                .fragmentDefinitionMap(fragmentDefinitionMap)
+                                .build();
+
+                        transformer.transform(visitor);
+
+                        this.eis.addAll(visitor.getExtractedEIs());
+                    });
                 }
                 catch (Exception ex) {
                     sink.error(ex);

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

@@ -0,0 +1,98 @@
+package com.intuit.graphql.orchestrator.utils;
+
+import graphql.language.Directive;
+import graphql.language.DirectivesContainer;
+import graphql.language.Field;
+import graphql.language.FragmentDefinition;
+import graphql.language.FragmentSpread;
+import graphql.language.InlineFragment;
+import graphql.language.Node;
+import graphql.language.OperationDefinition;
+import graphql.language.Selection;
+import graphql.language.SelectionSet;
+import org.apache.commons.lang3.ObjectUtils;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+public class NodeUtils {
+
+    /**
+     * Transforms node by removing desired directive from the node.
+     * Returns transformed node
+     * Throws an exception if it is at a location that is not supported.
+     * @param node: the selection that you would like to prune
+     * @param directiveName: the directive you would like to search for
+     * @return selection: pruned selection
+     * */
+    public static Selection removeDirectiveFromNode(Selection node, String directiveName) {
+        //DirectivesContainer returns getDirectives as a List<Object> so need to cast to stream correctly
+        List<Directive> prunedDirectives = ((List<Directive>)((DirectivesContainer)node).getDirectives())
+                .stream()
+                .filter(directive -> ObjectUtils.notEqual(directiveName, directive.getName()))
+                .collect(Collectors.toList());
+
+        if(node instanceof Field) {
+           return ((Field) node).transform(builder -> builder.directives(prunedDirectives));
+        } else if(node instanceof InlineFragment) {
+            return ((InlineFragment) node).transform(builder -> builder.directives(prunedDirectives));
+        }  else {
+            return ((FragmentSpread) node).transform(builder -> builder.directives(prunedDirectives));
+        }
+    }
+
+    /**
+     * Check if the selection has children selections
+     * @param node selection that you are trying to check
+     * @return boolean: true if selection doesn't have selectionset as child, false otherwise
+     * */
+    public static boolean isLeaf(Selection node) {
+        return node.getChildren().isEmpty() ||
+                node.getChildren()
+                .stream()
+                .noneMatch(SelectionSet.class::isInstance);
+    }
+
+
+    /**
+     * Generates new node
+     * Transforms the parentNode with a new selection set consisting of the pruned child and typename fields
+     * @param parentNode node that will be transformed and will contain only selections
+     * @param selections selections that need to be in selection set for parentNode
+     * @return node that only contains the passed in selections
+     * */
+    public static Node transformNodeWithSelections(Node parentNode, Selection... selections) {
+        SelectionSet prunedSelectionSet = SelectionSet.newSelectionSet()
+                .selections(Arrays.asList(selections))
+                .build();
+
+        if(parentNode instanceof Field) {
+            return ((Field) parentNode).transform(builder -> builder.selectionSet(prunedSelectionSet));
+        } else if (parentNode instanceof FragmentDefinition) {
+            //add fragment spread names here in case of nested fragment spreads
+            return ((FragmentDefinition) parentNode).transform(builder -> builder.selectionSet(prunedSelectionSet));
+        } else if (parentNode instanceof InlineFragment) {
+            return ((InlineFragment) parentNode).transform(builder -> builder.selectionSet(prunedSelectionSet));
+        } else {
+            return ((OperationDefinition) parentNode).transform(builder -> builder.selectionSet(prunedSelectionSet));
+        }
+    }
+
+    /**
+     * Retrieves the objects from map matching key
+     * @param kvMap: map that will be searched
+     * @param keyCollection keys to check
+     * @return list of values from map with matching key
+     * */
+    public static <K, V> List<V> getAllMapValuesWithMatchingKeys(Map<K, V> kvMap, Collection<K> keyCollection) {
+        return keyCollection
+                .stream()
+                .map(kvMap::get)
+                .filter(Objects::nonNull)
+                .collect(Collectors.toList());
+    }
+}

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

@@ -4,11 +4,12 @@
 import graphql.language.FragmentDefinition;
 import graphql.language.Node;
 import graphql.util.TraverserContext;
+import org.apache.commons.lang3.StringUtils;
+
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
-import org.apache.commons.lang3.StringUtils;
 
 public class QueryPathUtils {
 

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

@@ -0,0 +1,22 @@
+package com.intuit.graphql.orchestrator.utils;
+
+import graphql.language.Node;
+import graphql.language.SelectionSetContainer;
+import graphql.util.TraverserContext;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class TraverserContextUtils {
+    /**
+     * Returns the current nodes parent node definitions
+     * @param currentNode context for current node
+     * @return List of graphql definitions
+     * */
+    public static List getParentDefinitions(TraverserContext<Node> currentNode) {
+        return currentNode.getParentNodes()
+                .stream()
+                .filter(SelectionSetContainer.class::isInstance)
+                .collect(Collectors.toList());
+    }
+}