From 0949e5190c7c87f4d9ba7aad02fd941dbbdb9af9 Mon Sep 17 00:00:00 2001
From: Alexey Zhokhov <alexey@zhokhov.com>
Date: Tue, 24 Aug 2021 04:35:46 +0800
Subject: [PATCH] Added an ability to add headers before writing GraphQL
 response.

---
 .../servlet/HttpRequestInvokerImpl.java       |  5 ++-
 .../kickstart/servlet/ListenerHandler.java    |  4 +++
 .../servlet/core/GraphQLServletListener.java  | 33 +++++++++++++++++++
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/HttpRequestInvokerImpl.java b/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/HttpRequestInvokerImpl.java
index 9d10e398..829bf091 100644
--- a/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/HttpRequestInvokerImpl.java
+++ b/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/HttpRequestInvokerImpl.java
@@ -91,7 +91,10 @@ private CompletableFuture<Void> handle(
     return futureResult
         .thenApplyQueryResult()
         .thenAccept(
-            it -> writeResultResponse(futureResult.getInvocationInput(), it, request, response))
+            it -> {
+              listenerHandler.beforeFlush();
+              writeResultResponse(futureResult.getInvocationInput(), it, request, response);
+            })
         .thenAccept(it -> listenerHandler.onSuccess())
         .exceptionally(
             t ->
diff --git a/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/ListenerHandler.java b/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/ListenerHandler.java
index d9461002..faff966d 100644
--- a/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/ListenerHandler.java
+++ b/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/ListenerHandler.java
@@ -60,6 +60,10 @@ void runCallbacks(Consumer<RequestCallback> action) {
         });
   }
 
+  void beforeFlush() {
+    runCallbacks(it -> it.beforeFlush(request, response));
+  }
+
   void onSuccess() {
     runCallbacks(it -> it.onSuccess(request, response));
   }
diff --git a/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/core/GraphQLServletListener.java b/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/core/GraphQLServletListener.java
index 685d1faf..fc0feb34 100644
--- a/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/core/GraphQLServletListener.java
+++ b/graphql-java-servlet/src/main/java/graphql/kickstart/servlet/core/GraphQLServletListener.java
@@ -6,17 +6,50 @@
 /** @author Andrew Potter */
 public interface GraphQLServletListener {
 
+  /**
+   * Called this method when the request started processing.
+   * @param request http request
+   * @param response http response
+   * @return request callback or {@literal null}
+   */
   default RequestCallback onRequest(HttpServletRequest request, HttpServletResponse response) {
     return null;
   }
 
+  /**
+   * The callback which used to add additional listeners for GraphQL request execution.
+   */
   interface RequestCallback {
 
+    /**
+     * Called right before the response will be written and flushed. Can be used for applying some
+     * changes to the response object, like adding response headers.
+     * @param request http request
+     * @param response http response
+     */
+    default void beforeFlush(HttpServletRequest request, HttpServletResponse response) {}
+
+    /**
+     * Called when GraphQL invoked successfully and the response was written already.
+     * @param request http request
+     * @param response http response
+     */
     default void onSuccess(HttpServletRequest request, HttpServletResponse response) {}
 
+    /**
+     * Called when GraphQL was failed and the response was written already.
+     * @param request http request
+     * @param response http response
+     */
     default void onError(
         HttpServletRequest request, HttpServletResponse response, Throwable throwable) {}
 
+    /**
+     * Called finally once on both success and failed GraphQL invocation. The response is also
+     * already written.
+     * @param request http request
+     * @param response http response
+     */
     default void onFinally(HttpServletRequest request, HttpServletResponse response) {}
   }
 }