Merge pull request #149 from jmccaull/batch_handling

oliemansm · web-flow · commit 0217361483e2 · 2019-02-26T08:59:26.000+01:00
Upgrades to customize batch handling
diff --git a/src/main/java/graphql/servlet/AbstractGraphQLHttpServlet.java b/src/main/java/graphql/servlet/AbstractGraphQLHttpServlet.java
@@ -73,6 +73,12 @@ public abstract class AbstractGraphQLHttpServlet extends HttpServlet implements
     @Deprecated
     protected abstract GraphQLObjectMapper getGraphQLObjectMapper();
 
+    /**
+     * @deprecated override {@link #getConfiguration()} instead
+     */
+    @Deprecated
+    protected abstract GraphQLBatchExecutionHandlerFactory getBatchExecutionHandlerFactory();
+
     /**
      * @deprecated override {@link #getConfiguration()} instead
      */
@@ -85,6 +91,7 @@ protected GraphQLConfiguration getConfiguration() {
                 .with(getGraphQLObjectMapper())
                 .with(isAsyncServletMode())
                 .with(listeners)
+                .with(getBatchExecutionHandlerFactory())
                 .build();
     }
 
@@ -113,6 +120,7 @@ public void init(ServletConfig servletConfig) {
             GraphQLInvocationInputFactory invocationInputFactory = configuration.getInvocationInputFactory();
             GraphQLObjectMapper graphQLObjectMapper = configuration.getObjectMapper();
             GraphQLQueryInvoker queryInvoker = configuration.getQueryInvoker();
+            GraphQLBatchExecutionHandlerFactory batchExecutionHandlerFactory = configuration.getBatchExecutionHandlerFactory();
 
             String path = request.getPathInfo();
             if (path == null) {
@@ -125,7 +133,7 @@ public void init(ServletConfig servletConfig) {
                 if (query != null) {
 
                     if (isBatchedQuery(query)) {
-                        queryBatched(queryInvoker, graphQLObjectMapper, invocationInputFactory.createReadOnly(graphQLObjectMapper.readBatchedGraphQLRequest(query), request, response), response);
+                        queryBatched(batchExecutionHandlerFactory, queryInvoker, graphQLObjectMapper, invocationInputFactory.createReadOnly(graphQLObjectMapper.readBatchedGraphQLRequest(query), request, response), response);
                     } else {
                         final Map<String, Object> variables = new HashMap<>();
                         if (request.getParameter("variables") != null) {
@@ -147,6 +155,7 @@ public void init(ServletConfig servletConfig) {
             GraphQLInvocationInputFactory invocationInputFactory = configuration.getInvocationInputFactory();
             GraphQLObjectMapper graphQLObjectMapper = configuration.getObjectMapper();
             GraphQLQueryInvoker queryInvoker = configuration.getQueryInvoker();
+            GraphQLBatchExecutionHandlerFactory batchExecutionHandlerFactory = configuration.getBatchExecutionHandlerFactory();
 
             try {
                 if (APPLICATION_GRAPHQL.equals(request.getContentType())) {
@@ -181,7 +190,7 @@ public void init(ServletConfig servletConfig) {
                             GraphQLBatchedInvocationInput invocationInput =
                                     invocationInputFactory.create(graphQLRequests, request, response);
                             invocationInput.getContext().setParts(fileItems);
-                            queryBatched(queryInvoker, graphQLObjectMapper, invocationInput, response);
+                            queryBatched(batchExecutionHandlerFactory, queryInvoker, graphQLObjectMapper, invocationInput, response);
                             return;
                         } else {
                             GraphQLRequest graphQLRequest;
@@ -207,7 +216,7 @@ public void init(ServletConfig servletConfig) {
                     InputStream inputStream = asMarkableInputStream(request.getInputStream());
 
                     if (isBatchedQuery(inputStream)) {
-                        queryBatched(queryInvoker, graphQLObjectMapper, invocationInputFactory.create(graphQLObjectMapper.readBatchedGraphQLRequest(inputStream), request, response), response);
+                        queryBatched(batchExecutionHandlerFactory, queryInvoker, graphQLObjectMapper, invocationInputFactory.create(graphQLObjectMapper.readBatchedGraphQLRequest(inputStream), request, response), response);
                     } else {
                         query(queryInvoker, graphQLObjectMapper, invocationInputFactory.create(graphQLObjectMapper.readGraphQLRequest(inputStream), request, response), response);
                     }
@@ -364,21 +373,13 @@ private void query(GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper graphQL
         }
     }
 
-    private void queryBatched(GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper graphQLObjectMapper, GraphQLBatchedInvocationInput invocationInput, HttpServletResponse resp) throws Exception {
+    private void queryBatched(GraphQLBatchExecutionHandlerFactory batchInputHandlerFactory, GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper graphQLObjectMapper, GraphQLBatchedInvocationInput invocationInput, HttpServletResponse resp) throws Exception {
         resp.setContentType(APPLICATION_JSON_UTF8);
         resp.setStatus(STATUS_OK);
 
         Writer respWriter = resp.getWriter();
-        respWriter.write('[');
-
-        queryInvoker.query(invocationInput, (result, hasNext) -> {
-            respWriter.write(graphQLObjectMapper.serializeResultAsJson(result));
-            if (hasNext) {
-                respWriter.write(',');
-            }
-        });
 
-        respWriter.write(']');
+        queryInvoker.query(invocationInput, batchInputHandlerFactory.getBatchHandler(respWriter, graphQLObjectMapper));
     }
 
     private <R> List<R> runListeners(Function<? super GraphQLServletListener, R> action) {
diff --git a/src/main/java/graphql/servlet/BatchExecutionHandler.java b/src/main/java/graphql/servlet/BatchExecutionHandler.java
@@ -0,0 +1,19 @@
+package graphql.servlet;
+
+import graphql.ExecutionInput;
+import graphql.ExecutionResult;
+
+import java.util.function.BiFunction;
+
+/**
+ * @author Andrew Potter
+ */
+public interface BatchExecutionHandler {
+    /**
+     * Allows separating the logic of handling batch queries from how each individual query is resolved.
+     * @param batchedInvocationInput the batch query input
+     * @param queryFunction Function to produce query results.
+     */
+    void handleBatch(GraphQLBatchedInvocationInput batchedInvocationInput, BiFunction<GraphQLInvocationInput, ExecutionInput, ExecutionResult> queryFunction);
+}
+
diff --git a/src/main/java/graphql/servlet/DefaultGraphQLBatchExecutionHandlerFactory.java b/src/main/java/graphql/servlet/DefaultGraphQLBatchExecutionHandlerFactory.java
@@ -0,0 +1,47 @@
+package graphql.servlet;
+
+import graphql.ExecutionInput;
+import graphql.ExecutionResult;
+
+import java.io.IOException;
+import java.io.Writer;
+import java.util.Iterator;
+import java.util.function.BiFunction;
+
+public class DefaultGraphQLBatchExecutionHandlerFactory implements GraphQLBatchExecutionHandlerFactory {
+    @Override
+    public BatchExecutionHandler getBatchHandler(Writer respWriter, GraphQLObjectMapper graphQLObjectMapper) {
+        return new DefaultGraphQLBatchExecutionHandler(respWriter, graphQLObjectMapper);
+    }
+
+    private class DefaultGraphQLBatchExecutionHandler implements BatchExecutionHandler {
+
+        private final Writer respWriter;
+
+        private final GraphQLObjectMapper graphQLObjectMapper;
+
+        private DefaultGraphQLBatchExecutionHandler(Writer respWriter, GraphQLObjectMapper graphQLObjectMapper) {
+            this.respWriter = respWriter;
+            this.graphQLObjectMapper = graphQLObjectMapper;
+        }
+
+        @Override
+        public void handleBatch(GraphQLBatchedInvocationInput batchedInvocationInput, BiFunction<GraphQLInvocationInput, ExecutionInput,
+            ExecutionResult> queryFunction) {
+            Iterator<ExecutionInput> executionInputIterator = batchedInvocationInput.getExecutionInputs().iterator();
+            try {
+                respWriter.write("[");
+                while (executionInputIterator.hasNext()) {
+                    ExecutionResult result = queryFunction.apply(batchedInvocationInput, executionInputIterator.next());
+                    respWriter.write(graphQLObjectMapper.serializeResultAsJson(result));
+                    if (executionInputIterator.hasNext()) {
+                       respWriter.write(",");
+                    }
+                }
+                respWriter.write("]");
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+}
diff --git a/src/main/java/graphql/servlet/DefaultGraphQLServlet.java b/src/main/java/graphql/servlet/DefaultGraphQLServlet.java
@@ -25,6 +25,11 @@ protected GraphQLObjectMapper getGraphQLObjectMapper() {
         return GraphQLObjectMapper.newBuilder().build();
     }
 
+    @Override
+    protected GraphQLBatchExecutionHandlerFactory getBatchExecutionHandlerFactory() {
+        return null;
+    }
+
     @Override
     protected boolean isAsyncServletMode() {
         return false;
diff --git a/src/main/java/graphql/servlet/GraphQLBatchExecutionHandlerFactory.java b/src/main/java/graphql/servlet/GraphQLBatchExecutionHandlerFactory.java
@@ -0,0 +1,16 @@
+package graphql.servlet;
+
+import java.io.Writer;
+
+/**
+ * Interface to allow customization of how batched queries are handled.
+ */
+public interface GraphQLBatchExecutionHandlerFactory {
+    /**
+     * Produces an BatchExecutionHandler instance for a specific response. Can maintain state across each request within a batch.
+     * @param respWriter to send the response back
+     * @param graphQLObjectMapper to serialize results.
+     * @return a handler instance
+     */
+    BatchExecutionHandler getBatchHandler(Writer respWriter, GraphQLObjectMapper graphQLObjectMapper);
+}
diff --git a/src/main/java/graphql/servlet/GraphQLConfiguration.java b/src/main/java/graphql/servlet/GraphQLConfiguration.java
@@ -13,6 +13,7 @@ public class GraphQLConfiguration {
     private GraphQLInvocationInputFactory invocationInputFactory;
     private GraphQLQueryInvoker queryInvoker;
     private GraphQLObjectMapper objectMapper;
+    private GraphQLBatchExecutionHandlerFactory batchExecutionHandlerFactory;
     private List<GraphQLServletListener> listeners;
     private boolean asyncServletModeEnabled;
     private Executor asyncExecutor;
@@ -30,10 +31,11 @@ public static GraphQLConfiguration.Builder with(GraphQLInvocationInputFactory in
         return new Builder(invocationInputFactory);
     }
 
-    private GraphQLConfiguration(GraphQLInvocationInputFactory invocationInputFactory, GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper objectMapper, List<GraphQLServletListener> listeners, boolean asyncServletModeEnabled, Executor asyncExecutor, long subscriptionTimeout) {
+    private GraphQLConfiguration(GraphQLInvocationInputFactory invocationInputFactory, GraphQLQueryInvoker queryInvoker, GraphQLObjectMapper objectMapper, GraphQLBatchExecutionHandlerFactory batchExecutionHandlerFactory, List<GraphQLServletListener> listeners, boolean asyncServletModeEnabled, Executor asyncExecutor, long subscriptionTimeout) {
         this.invocationInputFactory = invocationInputFactory;
         this.queryInvoker = queryInvoker;
         this.objectMapper = objectMapper;
+        this.batchExecutionHandlerFactory = batchExecutionHandlerFactory;
         this.listeners = listeners;
         this.asyncServletModeEnabled = asyncServletModeEnabled;
         this.asyncExecutor = asyncExecutor;
@@ -52,6 +54,10 @@ public GraphQLObjectMapper getObjectMapper() {
         return objectMapper;
     }
 
+    public GraphQLBatchExecutionHandlerFactory getBatchExecutionHandlerFactory() {
+        return batchExecutionHandlerFactory;
+    }
+
     public List<GraphQLServletListener> getListeners() {
         return new ArrayList<>(listeners);
     }
@@ -82,6 +88,7 @@ public static class Builder {
         private GraphQLInvocationInputFactory invocationInputFactory;
         private GraphQLQueryInvoker queryInvoker = GraphQLQueryInvoker.newBuilder().build();
         private GraphQLObjectMapper objectMapper = GraphQLObjectMapper.newBuilder().build();
+        private GraphQLBatchExecutionHandlerFactory graphQLBatchExecutionHandlerFactory = new DefaultGraphQLBatchExecutionHandlerFactory();
         private List<GraphQLServletListener> listeners = new ArrayList<>();
         private boolean asyncServletModeEnabled = false;
         private Executor asyncExecutor = Executors.newCachedThreadPool(new GraphQLThreadFactory());
@@ -109,6 +116,13 @@ public Builder with(GraphQLObjectMapper objectMapper) {
             return this;
         }
 
+        public Builder with(GraphQLBatchExecutionHandlerFactory batchExecutionHandlerFactory) {
+            if (batchExecutionHandlerFactory != null) {
+                this.graphQLBatchExecutionHandlerFactory = batchExecutionHandlerFactory;
+            }
+            return this;
+        }
+
         public Builder with(List<GraphQLServletListener> listeners) {
             if (listeners != null) {
                 this.listeners = listeners;
@@ -148,6 +162,7 @@ public GraphQLConfiguration build() {
                     this.invocationInputFactory != null ? this.invocationInputFactory : invocationInputFactoryBuilder.build(),
                     queryInvoker,
                     objectMapper,
+                graphQLBatchExecutionHandlerFactory,
                     listeners,
                     asyncServletModeEnabled,
                     asyncExecutor,
diff --git a/src/main/java/graphql/servlet/GraphQLHttpServlet.java b/src/main/java/graphql/servlet/GraphQLHttpServlet.java
@@ -2,10 +2,6 @@
 
 import graphql.schema.GraphQLSchema;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Objects;
-
 /**
  * @author Michiel Oliemans
  */
@@ -37,6 +33,11 @@ protected GraphQLObjectMapper getGraphQLObjectMapper() {
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    protected GraphQLBatchExecutionHandlerFactory getBatchExecutionHandlerFactory() {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     protected boolean isAsyncServletMode() {
         throw new UnsupportedOperationException();
diff --git a/src/main/java/graphql/servlet/GraphQLQueryInvoker.java b/src/main/java/graphql/servlet/GraphQLQueryInvoker.java
@@ -11,13 +11,11 @@
 import graphql.execution.preparsed.NoOpPreparsedDocumentProvider;
 import graphql.execution.preparsed.PreparsedDocumentProvider;
 import graphql.schema.GraphQLSchema;
-import graphql.servlet.internal.ExecutionResultHandler;
 
 import javax.security.auth.Subject;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 import java.util.ArrayList;
-import java.util.Iterator;
 import java.util.List;
 import java.util.function.Supplier;
 
@@ -42,13 +40,8 @@ public ExecutionResult query(GraphQLSingleInvocationInput singleInvocationInput)
         return query(singleInvocationInput, singleInvocationInput.getExecutionInput());
     }
 
-    public void query(GraphQLBatchedInvocationInput batchedInvocationInput, ExecutionResultHandler executionResultHandler) {
-        Iterator<ExecutionInput> executionInputIterator = batchedInvocationInput.getExecutionInputs().iterator();
-
-        while (executionInputIterator.hasNext()) {
-            ExecutionResult result = query(batchedInvocationInput, executionInputIterator.next());
-            executionResultHandler.accept(result, executionInputIterator.hasNext());
-        }
+    public void query(GraphQLBatchedInvocationInput batchedInvocationInput, BatchExecutionHandler batchExecutionHandler) {
+        batchExecutionHandler.handleBatch(batchedInvocationInput, this::query);
     }
 
     private GraphQL newGraphQL(GraphQLSchema schema, Object context) {
diff --git a/src/main/java/graphql/servlet/OsgiGraphQLHttpServlet.java b/src/main/java/graphql/servlet/OsgiGraphQLHttpServlet.java
@@ -46,6 +46,7 @@ public class OsgiGraphQLHttpServlet extends AbstractGraphQLHttpServlet {
     private InstrumentationProvider instrumentationProvider = new NoOpInstrumentationProvider();
     private GraphQLErrorHandler errorHandler = new DefaultGraphQLErrorHandler();
     private PreparsedDocumentProvider preparsedDocumentProvider = NoOpPreparsedDocumentProvider.INSTANCE;
+    private GraphQLBatchExecutionHandlerFactory executionResultHandlerFactory = new DefaultGraphQLBatchExecutionHandlerFactory();
 
     private GraphQLSchemaProvider schemaProvider;
 
@@ -84,6 +85,11 @@ protected GraphQLObjectMapper getGraphQLObjectMapper() {
         return graphQLObjectMapper;
     }
 
+    @Override
+    protected GraphQLBatchExecutionHandlerFactory getBatchExecutionHandlerFactory() {
+        return executionResultHandlerFactory;
+    }
+
     @Override
     protected boolean isAsyncServletMode() {
         return false;
diff --git a/src/main/java/graphql/servlet/SimpleGraphQLHttpServlet.java b/src/main/java/graphql/servlet/SimpleGraphQLHttpServlet.java
@@ -69,6 +69,11 @@ protected GraphQLObjectMapper getGraphQLObjectMapper() {
         return configuration.getObjectMapper();
     }
 
+    @Override
+    protected GraphQLBatchExecutionHandlerFactory getBatchExecutionHandlerFactory() {
+        return configuration.getBatchExecutionHandlerFactory();
+    }
+
     @Override
     protected boolean isAsyncServletMode() {
         return configuration.isAsyncServletModeEnabled();
diff --git a/src/main/java/graphql/servlet/internal/ExecutionResultHandler.java b/src/main/java/graphql/servlet/internal/ExecutionResultHandler.java
diff --git a/src/test/groovy/graphql/servlet/AbstractGraphQLHttpServletSpec.groovy b/src/test/groovy/graphql/servlet/AbstractGraphQLHttpServletSpec.groovy
diff --git a/src/test/groovy/graphql/servlet/TestBatchExecutionHandlerFactoryBatch.java b/src/test/groovy/graphql/servlet/TestBatchExecutionHandlerFactoryBatch.java
diff --git a/src/test/groovy/graphql/servlet/TestUtils.groovy b/src/test/groovy/graphql/servlet/TestUtils.groovy
