rewrote execution result handler to allow more customization of batch query processing

Author: jmccaull

src/main/java/graphql/servlet/DefaultGraphQLExecutionResultHandlerFactory.java

@@ -1,12 +1,12 @@
 package graphql.servlet;
 
+import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 
 import java.io.IOException;
 import java.io.Writer;
-import java.util.ArrayList;
 import java.util.Iterator;
-import java.util.List;
+import java.util.function.BiFunction;
 
 public class DefaultGraphQLExecutionResultHandlerFactory implements GraphQLExecutionResultHandlerFactory {
     @Override
@@ -16,8 +16,6 @@ public ExecutionResultHandler getBatchHandler(Writer respWriter, GraphQLObjectMa
 
     private class DefaultGraphQLExecutionResultHandler implements ExecutionResultHandler {
 
-        private final List<ExecutionResult> results = new ArrayList<>();
-
         private final Writer respWriter;
 
         private final GraphQLObjectMapper graphQLObjectMapper;
@@ -28,19 +26,16 @@ private DefaultGraphQLExecutionResultHandler(Writer respWriter, GraphQLObjectMap
         }
 
         @Override
-        public void accept(ExecutionResult result) {
-            results.add(result);
-        }
-
-        @Override
-        public void finalizeResults() {
+        public void handleBatch(GraphQLBatchedInvocationInput batchedInvocationInput, BiFunction<GraphQLInvocationInput, ExecutionInput,
+            ExecutionResult> queryFunction) {
+            Iterator<ExecutionInput> executionInputIterator = batchedInvocationInput.getExecutionInputs().iterator();
             try {
                 respWriter.write("[");
-                Iterator<ExecutionResult> iterator = results.iterator();
-                while (iterator.hasNext()) {
-                    respWriter.write(graphQLObjectMapper.serializeResultAsJson(iterator.next()));
-                    if (iterator.hasNext()) {
-                        respWriter.write(",");
+                while (executionInputIterator.hasNext()) {
+                    ExecutionResult result = queryFunction.apply(batchedInvocationInput, executionInputIterator.next());
+                    respWriter.write(graphQLObjectMapper.serializeResultAsJson(result));
+                    if (executionInputIterator.hasNext()) {
+                       respWriter.write(",");
                     }
                 }
                 respWriter.write("]");

src/main/java/graphql/servlet/ExecutionResultHandler.java

@@ -4,24 +4,13 @@
 import graphql.ExecutionResult;
 
 import java.util.Iterator;
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
 
 /**
  * @author Andrew Potter
  */
-public interface ExecutionResultHandler extends Consumer<ExecutionResult> {
-    /**
-     * Allows the number of queries in a batch to be limited.
-     * @param executionInputIterator iterator for the current set of requests.
-     * @return if processing should continue.
-     */
-    default boolean shouldContinue(Iterator<ExecutionInput> executionInputIterator){
-        return executionInputIterator.hasNext();
-    }
-
-    /**
-     * Should perform any actions, such as formatting and sending results, to complete a batch.
-     */
-    void finalizeResults();
+public interface ExecutionResultHandler {
+    void handleBatch(GraphQLBatchedInvocationInput batchedInvocationInput, BiFunction<GraphQLInvocationInput, ExecutionInput, ExecutionResult> queryFunction);
 }
 

src/main/java/graphql/servlet/GraphQLQueryInvoker.java

@@ -42,14 +42,7 @@ public ExecutionResult query(GraphQLSingleInvocationInput singleInvocationInput)
     }
 
     public void query(GraphQLBatchedInvocationInput batchedInvocationInput, ExecutionResultHandler executionResultHandler) {
-        Iterator<ExecutionInput> executionInputIterator = batchedInvocationInput.getExecutionInputs().iterator();
-
-        while (executionResultHandler.shouldContinue(executionInputIterator)) {
-            ExecutionResult result = query(batchedInvocationInput, executionInputIterator.next());
-            executionResultHandler.accept(result);
-        }
-
-        executionResultHandler.finalizeResults();
+        executionResultHandler.handleBatch(batchedInvocationInput, this::query);
     }
 
     private GraphQL newGraphQL(GraphQLSchema schema, Object context) {

src/test/groovy/graphql/servlet/TestBatchInputHandlerFactory.java

@@ -5,9 +5,10 @@
 
 import java.io.IOException;
 import java.io.Writer;
-import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.function.BiFunction;
+import java.util.stream.Collectors;
 
 public class TestBatchInputHandlerFactory implements GraphQLExecutionResultHandlerFactory {
     @Override
@@ -19,15 +20,23 @@ private class LimitedBatchSizeHandler implements ExecutionResultHandler {
 
         Writer writer;
         GraphQLObjectMapper mapper;
-        List<ExecutionResult> results = new ArrayList<>();
 
         private LimitedBatchSizeHandler(Writer respWriter, GraphQLObjectMapper graphQLObjectMapper) {
             this.mapper = graphQLObjectMapper;
             this.writer = respWriter;
         }
 
         @Override
-        public void finalizeResults() {
+        public void handleBatch(GraphQLBatchedInvocationInput batchedInvocationInput, BiFunction<GraphQLInvocationInput, ExecutionInput,
+            ExecutionResult> queryFunction) {
+            List<ExecutionResult> results = batchedInvocationInput.getExecutionInputs().parallelStream()
+                .limit(2)
+                .map(input -> queryFunction.apply(batchedInvocationInput, input))
+                .collect(Collectors.toList());
+            writeResults(results);
+        }
+
+        public void writeResults(List<ExecutionResult> results) {
             try {
                 writer.write("[");
                 Iterator<ExecutionResult> iter = results.iterator();
@@ -42,15 +51,5 @@ public void finalizeResults() {
                 throw new RuntimeException(e);
             }
         }
-
-        @Override
-        public void accept(ExecutionResult executionResult) {
-            results.add(executionResult);
-        }
-
-        @Override
-        public boolean shouldContinue(Iterator<ExecutionInput> executionInputIterator) {
-            return results.size() < 2 && executionInputIterator.hasNext();
-        }
     }
 }