Defer support with RxJava

Author: JacobMountain

build.gradle

@@ -53,6 +53,7 @@ dependencies {
     // Servlet
     compile 'javax.servlet:javax.servlet-api:3.1.0'
     compile 'javax.websocket:javax.websocket-api:1.1'
+    compile 'io.projectreactor:reactor-core:3.0.5.RELEASE'
 
     // GraphQL
     compile "com.graphql-java:graphql-java:$LIB_GRAPHQL_JAVA_VER"

src/main/java/graphql/servlet/AbstractGraphQLHttpServlet.java

@@ -2,7 +2,10 @@
 
 import com.google.common.io.ByteStreams;
 import com.google.common.io.CharStreams;
+import graphql.DeferredExecutionResult;
 import graphql.ExecutionResult;
+import graphql.GraphQL;
+import graphql.execution.reactive.SingleSubscriberPublisher;
 import graphql.introspection.IntrospectionQuery;
 import graphql.schema.GraphQLFieldDefinition;
 import graphql.servlet.config.GraphQLConfiguration;
@@ -23,6 +26,7 @@
 import org.reactivestreams.Subscription;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import reactor.core.publisher.Flux;
 
 import javax.servlet.AsyncContext;
 import javax.servlet.AsyncEvent;
@@ -373,7 +377,9 @@ private void query(GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper graphQL
                        HttpServletRequest req, HttpServletResponse resp) throws IOException {
         ExecutionResult result = queryInvoker.query(invocationInput);
 
-        if (!(result.getData() instanceof Publisher)) {
+        boolean isDeferred = Objects.nonNull(result.getExtensions()) && result.getExtensions().containsKey(GraphQL.DEFERRED_RESULTS);
+
+        if (!(result.getData() instanceof Publisher || isDeferred)) {
             resp.setContentType(APPLICATION_JSON_UTF8);
             resp.setStatus(STATUS_OK);
             resp.getWriter().write(graphQLObjectMapper.serializeResultAsJson(result));
@@ -390,7 +396,15 @@ private void query(GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper graphQL
             AtomicReference<Subscription> subscriptionRef = new AtomicReference<>();
             asyncContext.addListener(new SubscriptionAsyncListener(subscriptionRef));
             ExecutionResultSubscriber subscriber = new ExecutionResultSubscriber(subscriptionRef, asyncContext, graphQLObjectMapper);
-            ((Publisher<ExecutionResult>) result.getData()).subscribe(subscriber);
+            Publisher<ExecutionResult> publisher;
+            if (result.getData() instanceof Publisher) {
+                publisher = result.getData();
+            } else {
+                publisher = new SingleSubscriberPublisher<>();
+                ((SingleSubscriberPublisher<ExecutionResult>) publisher).offer(result);
+                publisher = Flux.merge(publisher, (Publisher<DeferredExecutionResult>) result.getExtensions().get(GraphQL.DEFERRED_RESULTS));
+            }
+            publisher.subscribe(subscriber);
             if (isInAsyncThread) {
                 // We need to delay the completion of async context until after the subscription has terminated, otherwise the AsyncContext is prematurely closed.
                 try {

src/main/java/graphql/servlet/core/GraphQLObjectMapper.java

@@ -5,9 +5,8 @@
 import com.fasterxml.jackson.databind.MappingIterator;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.ObjectReader;
-import graphql.ExecutionResult;
-import graphql.ExecutionResultImpl;
-import graphql.GraphQLError;
+import graphql.*;
+import graphql.execution.ExecutionPath;
 import graphql.servlet.config.ConfiguringObjectMapperProvider;
 import graphql.servlet.config.ObjectMapperConfigurer;
 import graphql.servlet.config.ObjectMapperProvider;
@@ -117,12 +116,19 @@ public ExecutionResult sanitizeErrors(ExecutionResult executionResult) {
         } else {
             errors = null;
         }
-
         return new ExecutionResultImpl(data, errors, extensions);
     }
 
     public Map<String, Object> createResultFromExecutionResult(ExecutionResult executionResult) {
-        return convertSanitizedExecutionResult(sanitizeErrors(executionResult));
+        ExecutionResult sanitizedExecutionResult = sanitizeErrors(executionResult);
+        if (executionResult instanceof DeferredExecutionResult) {
+            sanitizedExecutionResult = DeferredExecutionResultImpl
+                    .newDeferredExecutionResult()
+                    .from(executionResult)
+                    .path(ExecutionPath.fromList(((DeferredExecutionResult) executionResult).getPath()))
+                    .build();
+        }
+        return convertSanitizedExecutionResult(sanitizedExecutionResult);
     }
 
     public Map<String, Object> convertSanitizedExecutionResult(ExecutionResult executionResult) {
@@ -144,6 +150,10 @@ public Map<String, Object> convertSanitizedExecutionResult(ExecutionResult execu
             result.put("data", executionResult.getData());
         }
 
+        if (executionResult instanceof DeferredExecutionResult) {
+            result.put("path", ((DeferredExecutionResult) executionResult).getPath());
+        }
+
         return result;
     }
 

src/test/groovy/graphql/servlet/AbstractGraphQLHttpServletSpec.groovy

@@ -283,6 +283,28 @@ class AbstractGraphQLHttpServletSpec extends Specification {
         getBatchedResponseContent()[1].data.echo == "test"
     }
 
+
+    def "deferred query over HTTP GET"() {
+        setup:
+        request.addParameter('query', 'query { deferred(arg:"test") @defer }')
+
+        when:
+        servlet.doGet(request, response)
+
+        then:
+        response.getStatus() == STATUS_OK
+        response.getContentType() == CONTENT_TYPE_SERVER_SENT_EVENTS
+        getSubscriptionResponseContent()[0].data.deferred == null
+
+        when:
+        subscriptionLatch.await(1, TimeUnit.SECONDS)
+
+        then:
+        def content = getSubscriptionResponseContent()
+        content[1].data == "test"
+        content[1].path == ["deferred"]
+    }
+
     def "Batch Execution Handler allows limiting batches and sending error messages."() {
         setup:
         servlet = TestUtils.createBatchCustomizedServlet({ env -> env.arguments.arg }, { env -> env.arguments.arg }, { env ->
@@ -1030,6 +1052,31 @@ class AbstractGraphQLHttpServletSpec extends Specification {
         getSubscriptionResponseContent()[1].data.echo == "Second\n\ntest"
     }
 
+    def "defer query over HTTP POST"() {
+        setup:
+        servlet = TestUtils.createDefaultServlet()
+        request.setContent(mapper.writeValueAsBytes([
+                query: 'query { deferred(arg:"test") @defer }'
+        ]))
+        request.setAsyncSupported(true)
+
+        when:
+        servlet.doPost(request, response)
+
+        then:
+        response.getStatus() == STATUS_OK
+        response.getContentType() == CONTENT_TYPE_SERVER_SENT_EVENTS
+        getSubscriptionResponseContent()[0].data.deferred == null
+
+        when:
+        subscriptionLatch.await(1, TimeUnit.SECONDS)
+
+        then:
+        def content = getSubscriptionResponseContent()
+        content[1].data == "test"
+        content[1].path == ["deferred"]
+    }
+
     def "errors before graphql schema execution return internal server error"() {
         setup:
         servlet = SimpleGraphQLHttpServlet.newBuilder(GraphQLInvocationInputFactory.newBuilder {

src/test/groovy/graphql/servlet/TestUtils.groovy

@@ -1,6 +1,7 @@
 package graphql.servlet
 
 import com.google.common.io.ByteStreams
+import graphql.Directives
 import graphql.Scalars
 import graphql.execution.reactive.SingleSubscriberPublisher
 import graphql.schema.*
@@ -15,6 +16,7 @@ import graphql.servlet.core.ApolloScalars
 import graphql.servlet.input.BatchInputPreProcessor
 import graphql.servlet.context.ContextSetting
 
+import java.util.concurrent.CompletableFuture
 import java.util.concurrent.atomic.AtomicReference
 
 class TestUtils {
@@ -95,7 +97,7 @@ class TestUtils {
     static def createGraphQlSchema(DataFetcher queryDataFetcher = { env -> env.arguments.arg },
                                    DataFetcher mutationDataFetcher = { env -> env.arguments.arg },
                                    DataFetcher subscriptionDataFetcher = { env ->
-                                       AtomicReference<SingleSubscriberPublisher<String>> publisherRef = new AtomicReference<>();
+                                       AtomicReference<SingleSubscriberPublisher<String>> publisherRef = new AtomicReference<>()
                                        publisherRef.set(new SingleSubscriberPublisher<>({ subscription ->
                                            publisherRef.get().offer(env.arguments.arg)
                                            publisherRef.get().noMoreData()
@@ -118,6 +120,20 @@ class TestUtils {
             field.type(new GraphQLNonNull(Scalars.GraphQLString))
             field.dataFetcher({ env -> null })
         }
+        .field { GraphQLFieldDefinition.Builder field ->
+            field.name("deferred")
+            field.type(Scalars.GraphQLString)
+            field.argument { argument ->
+                argument.name("arg")
+                argument.type(Scalars.GraphQLString)
+            }
+            field.dataFetcher({ env ->
+                return CompletableFuture.supplyAsync( {
+                    Thread.sleep(100)
+                    env.arguments.arg
+                })
+            })
+        }
         .build()
 
         GraphQLObjectType mutation = GraphQLObjectType.newObject()
@@ -174,6 +190,7 @@ class TestUtils {
                 .mutation(mutation)
                 .subscription(subscription)
                 .additionalType(ApolloScalars.Upload)
+                .additionalDirective(Directives.DeferDirective)
                 .build()
     }