Add suport for defer

pom.xml

@@ -2,7 +2,7 @@
   <modelVersion>4.0.0</modelVersion>
   <groupId>com.intuit.graphql</groupId>
   <artifactId>graphql-orchestrator-java</artifactId>
-  <version>5.0.19-SNAPSHOT</version>
+  <version>5.0.20-DEFER-SNAPSHOT</version>
   <packaging>jar</packaging>
   <name>graphql-orchestrator-java</name>
   <description>GraphQL Orchestrator combines multiple graphql services into a single unified schema</description>
@@ -39,6 +39,7 @@
     <graphql-sdl-version>3.0.1</graphql-sdl-version>
     <xtextVersion>2.26.0</xtextVersion>
     <graphQLVersion>17.4</graphQLVersion>
+    <reactorVersion>3.4.24</reactorVersion>
   </properties>
 
   <profiles>
@@ -159,6 +160,17 @@
       <artifactId>jackson-databind</artifactId>
       <version>2.13.4.1</version>
     </dependency>
+    <dependency>
+        <groupId>io.projectreactor</groupId>
+        <artifactId>reactor-core</artifactId>
+        <version>${reactorVersion}</version>
+    </dependency>
+    <dependency>
+      <groupId>io.projectreactor</groupId>
+      <artifactId>reactor-test</artifactId>
+      <version>${reactorVersion}</version>
+      <scope>test</scope>
+    </dependency>
 
     <!-- TEST DEPENDENCIES -->
     <dependency>

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

@@ -1,8 +1,12 @@
 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;
 import graphql.ExecutionResult;
+import graphql.ExecutionResultImpl;
 import graphql.GraphQL;
 import graphql.GraphQLContext;
 import graphql.execution.AsyncExecutionStrategy;
@@ -11,28 +15,35 @@
 import graphql.execution.instrumentation.ChainedInstrumentation;
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation;
+import graphql.execution.reactive.SubscriptionPublisher;
 import graphql.schema.GraphQLSchema;
 import lombok.extern.slf4j.Slf4j;
 import org.dataloader.BatchLoader;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderRegistry;
+import reactor.core.publisher.Flux;
+import reactor.core.scheduler.Schedulers;
 
 import java.util.Arrays;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
 
+import static com.intuit.graphql.orchestrator.utils.DirectivesUtil.USE_DEFER;
 import static java.util.Objects.requireNonNull;
 
 @Slf4j
 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;
@@ -41,8 +52,8 @@ public class GraphQLOrchestrator {
   private final ExecutionStrategy mutationExecutionStrategy;
 
   private GraphQLOrchestrator(final RuntimeGraph runtimeGraph, final List<Instrumentation> instrumentations,
-                              final ExecutionIdProvider executionIdProvider, final ExecutionStrategy queryExecutionStrategy,
-                              final ExecutionStrategy mutationExecutionStrategy) {
+      final ExecutionIdProvider executionIdProvider, final ExecutionStrategy queryExecutionStrategy,
+      final ExecutionStrategy mutationExecutionStrategy) {
     this.runtimeGraph = runtimeGraph;
     this.instrumentations = instrumentations;
     this.executionIdProvider = executionIdProvider;
@@ -63,16 +74,80 @@ private DataLoaderRegistry buildNewDataLoaderRegistry() {
     // to create a new DataLoader per request. Else it will use the cache which is shared
     // across request.
     final Map<BatchLoader, DataLoader> temporaryMap = this.runtimeGraph.getBatchLoaderMap().values().stream().distinct()
-            .collect(Collectors.toMap(Function.identity(), DataLoader::new));
+        .collect(Collectors.toMap(Function.identity(), DataLoader::new));
 
     this.runtimeGraph.getBatchLoaderMap()
-            .forEach((key, batchLoader) ->
-                    dataLoaderRegistry.register(key, temporaryMap.getOrDefault(batchLoader, new DataLoader(batchLoader))));
+        .forEach((key, batchLoader) ->
+            dataLoaderRegistry.register(key, temporaryMap.getOrDefault(batchLoader, new DataLoader(batchLoader))));
     return dataLoaderRegistry;
   }
 
   public CompletableFuture<ExecutionResult> execute(ExecutionInput executionInput) {
+    return execute(executionInput, DEFAULT_DEFER_OPTIONS, DISABLED_DEFER);
+  }
+
+  public CompletableFuture<ExecutionResult> execute(ExecutionInput executionInput, DeferOptions deferOptions, boolean hasDefer) {
+    if(hasDefer) {
+      return executeWithDefer(executionInput, deferOptions);
+    }
+      final GraphQL graphQL = constructGraphQL();
+
+      final ExecutionInput newExecutionInput = executionInput
+              .transform(builder -> builder.dataLoaderRegistry(buildNewDataLoaderRegistry()));
+
+      if (newExecutionInput.getContext() instanceof GraphQLContext) {
+        ((GraphQLContext) newExecutionInput.getContext())
+                .put(DATA_LOADER_REGISTRY_CONTEXT_KEY, newExecutionInput.getDataLoaderRegistry());
+        ((GraphQLContext) newExecutionInput.getContext())
+                .put(USE_DEFER , false);
+      }
+      return graphQL.executeAsync(newExecutionInput);
+  }
 
+  private CompletableFuture<ExecutionResult> executeWithDefer(ExecutionInput executionInput, DeferOptions options) {
+    AtomicInteger responses = new AtomicInteger(0);
+    MultiEIGenerator eiGenerator = new MultiEIGenerator(executionInput, options, this.getSchema());
+
+    Flux<Object> executionResultPublisher = eiGenerator.generateEIs()
+            .filter(ei -> !ei.getQuery().equals(""))
+            .publishOn(Schedulers.elastic())
+            .map(ei -> {
+              log.error("Timestamp processing emittedValue: {}", System.currentTimeMillis());
+              return this.generateEIWIthNewContext(ei);
+            })
+            .map(constructGraphQL()::executeAsync)
+            .map(CompletableFuture::join)
+            .doOnNext(executionResult -> responses.getAndIncrement())
+            .map(ExecutionResultImpl.newExecutionResult()::from)
+            .map(builder -> builder.addExtension("hasMoreData", hasMoreData(eiGenerator.getNumOfEIs(), responses.get())))
+            .map(ExecutionResultImpl.Builder::build)
+            .map(Object.class::cast)
+            .takeUntil(object -> eiGenerator.getNumOfEIs() != null && !hasMoreData(eiGenerator.getNumOfEIs(), responses.get()));
+
+    SubscriptionPublisher multiResultPublisher = new SubscriptionPublisher(executionResultPublisher, null);
+
+    return CompletableFuture.completedFuture(ExecutionResultImpl.newExecutionResult().data(multiResultPublisher).build());
+  }
+
+  private boolean hasMoreData(Integer expectedNumOfEIs, Integer numOfResponses) {
+    return expectedNumOfEIs == null || expectedNumOfEIs.intValue() != numOfResponses.intValue();
+  }
+  private ExecutionInput generateEIWIthNewContext(ExecutionInput ei) {
+    DataLoaderRegistry registry = buildNewDataLoaderRegistry();
+
+    GraphQLContext graphqlContext = GraphQLContext.newContext()
+            .of((GraphQLContext)ei.getContext())
+            .put(DATA_LOADER_REGISTRY_CONTEXT_KEY, registry)
+            .put(USE_DEFER, true)
+            .build();
+
+    return ei.transform(builder -> {
+      builder.dataLoaderRegistry(registry);
+      builder.context(graphqlContext);
+    });
+  }
+
+  private GraphQL constructGraphQL() {
     final GraphQL.Builder graphqlBuilder = GraphQL.newGraphQL(runtimeGraph.getExecutableSchema())
             .instrumentation(new ChainedInstrumentation(instrumentations))
             .executionIdProvider(executionIdProvider)
@@ -81,17 +156,7 @@ public CompletableFuture<ExecutionResult> execute(ExecutionInput executionInput)
     if (Objects.nonNull(mutationExecutionStrategy)) {
       graphqlBuilder.mutationExecutionStrategy(mutationExecutionStrategy);
     }
-
-    final GraphQL graphQL = graphqlBuilder.build();
-
-    final ExecutionInput newExecutionInput = executionInput
-            .transform(builder -> builder.dataLoaderRegistry(buildNewDataLoaderRegistry()));
-
-    if (newExecutionInput.getContext() instanceof GraphQLContext) {
-      ((GraphQLContext) executionInput.getContext())
-              .put(DATA_LOADER_REGISTRY_CONTEXT_KEY, newExecutionInput.getDataLoaderRegistry());
-    }
-    return graphQL.executeAsync(newExecutionInput);
+    return graphqlBuilder.build();
   }
 
   public GraphQLSchema getSchema() {
@@ -113,7 +178,7 @@ public static class Builder {
     private ExecutionStrategy queryExecutionStrategy = new AsyncExecutionStrategy();
     private ExecutionStrategy mutationExecutionStrategy = null;
     private List<Instrumentation> instrumentations = new LinkedList<>(
-            Arrays.asList(new DataLoaderDispatcherInstrumentation()));
+        Arrays.asList(new DataLoaderDispatcherInstrumentation(), new DeferDirectiveInstrumentation()));
 
     private Builder() {
     }
@@ -155,7 +220,7 @@ public Builder mutationExecutionStrategy(final ExecutionStrategy mutationExecuti
 
     public GraphQLOrchestrator build() {
       return new GraphQLOrchestrator(runtimeGraph, instrumentations, executionIdProvider, queryExecutionStrategy,
-              mutationExecutionStrategy);
+          mutationExecutionStrategy);
     }
   }
 }

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

@@ -3,20 +3,19 @@
 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;
 import graphql.schema.GraphQLType;
-import java.util.Map;
 import lombok.Builder;
 import lombok.NonNull;
 
-@Builder
-public class DownstreamQueryRedactor {
+import java.util.Map;
 
-  private static final AstTransformer AST_TRANSFORMER = new AstTransformer();
+import static com.intuit.graphql.orchestrator.utils.GraphQLUtil.AST_TRANSFORMER;
 
+@Builder
+public class DownstreamQueryRedactor {
   @NonNull private Node<?> root;
   @NonNull private GraphQLType rootType;
   @NonNull private GraphQLType rootParentType;
@@ -31,7 +30,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;
 }