feat: add ability to propagate HTTP headers with callbacks (#435)

Adds a mechanism for propgating HTTP headers from the original GraphQL
subscription request to the subsequent callback responses.

```java
@bean
public CallbackWebGraphQLInterceptor callbackGraphQlInterceptor(
    SubscriptionCallbackHandler callbackHandler) {
  var headers = Set.of("X-Custom-Header", "X-Custom-Header2");
  return new CallbackWebGraphQLInterceptor(callbackHandler, headers);
}
```

Author: dariuszkuc

spring-subscription-callback/README.md

@@ -93,6 +93,7 @@ public class GraphQLConfiguration {
     // to allow users to still apply other interceptors that handle common stuff (e.g. extracting
     // auth headers, etc).
     // You can override this behavior by specifying custom order.
+    // You can also provide a set of HTTP headers that should be propagated to the callback responses.
     @Bean
     public CallbackWebGraphQLInterceptor callbackGraphQlInterceptor(
             SubscriptionCallbackHandler callbackHandler) {
@@ -113,6 +114,8 @@ public class GraphQLConfiguration {
 }
 ```
 
+### Custom Scheduler
+
 By default, subscription and heartbeat stream will be executed in non-blocking way on [bounded elastic](https://projectreactor.io/docs/core/release/api/reactor/core/scheduler/Schedulers.html#boundedElastic--)
 scheduler. If you need more granular control over the underlying scheduler, you can configure callback handler to run on
 your provided scheduler.
@@ -124,3 +127,17 @@ public SubscriptionCallbackHandler callbackHandler(ExecutionGraphQlService graph
     return new SubscriptionCallbackHandler(graphQlService, customScheduler);
 }
 ```
+
+### Propagating HTTP Headers
+
+Subscription callback interceptor can be configured with a list of header names to propagate from the
+original GraphQL subscription request.
+
+```java
+@Bean
+public CallbackWebGraphQLInterceptor callbackGraphQlInterceptor(
+    SubscriptionCallbackHandler callbackHandler) {
+  var headers = Set.of("X-Custom-Header");
+  return new CallbackWebGraphQLInterceptor(callbackHandler, headers);
+}
+```

spring-subscription-callback/src/main/java/com/apollographql/subscription/CallbackWebGraphQLInterceptor.java

@@ -7,6 +7,11 @@
 import com.apollographql.subscription.callback.SubscriptionCallbackHandler;
 import graphql.ExecutionResult;
 import graphql.GraphQLError;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.springframework.core.Ordered;
@@ -20,13 +25,15 @@
 import reactor.core.publisher.Mono;
 
 /**
- * Interceptor that provides support for Apollo Subscription Callback Protocol. This interceptor
- * defaults to {@link Ordered#LOWEST_PRECEDENCE} order as it should run last in chain to allow users
- * to still apply other interceptors that handle common stuff (e.g. extracting auth headers, etc).
- * You can override this behavior by specifying custom order.
+ * Interceptor that provides support for Apollo Subscription Callback Protocol. HTTP headers from
+ * the original subscription GraphQL operation can be propagated to the subsequent callback
+ * responses by specifying a Set of contextual header names. This interceptor defaults to {@link
+ * Ordered#LOWEST_PRECEDENCE} order as it should run last in the chain to allow users to still apply
+ * other interceptors that handle common stuff (e.g. extracting auth headers, etc). You can override
+ * this behavior by specifying custom order.
  *
  * @see <a
- *     href="https://www.apollographql.com/docs/router/executing-operations/subscription-callback-protocol">Subscription
+ *     href="https://www.apollographql.com/docs/graphos/routing/operations/subscriptions/callback-protocol">Subscription
  *     Callback Protocol</a>
  */
 public class CallbackWebGraphQLInterceptor implements WebGraphQlInterceptor, Ordered {
@@ -35,15 +42,29 @@ public class CallbackWebGraphQLInterceptor implements WebGraphQlInterceptor, Ord
 
   private final SubscriptionCallbackHandler subscriptionCallbackHandler;
   private final int order;
+  private final Set<String> contextualHeaders;
 
   public CallbackWebGraphQLInterceptor(SubscriptionCallbackHandler subscriptionCallbackHandler) {
     this(subscriptionCallbackHandler, LOWEST_PRECEDENCE);
   }
 
   public CallbackWebGraphQLInterceptor(
       SubscriptionCallbackHandler subscriptionCallbackHandler, int order) {
+    this(subscriptionCallbackHandler, order, new HashSet<>());
+  }
+
+  public CallbackWebGraphQLInterceptor(
+      SubscriptionCallbackHandler subscriptionCallbackHandler, Set<String> contextualHeaders) {
+    this(subscriptionCallbackHandler, LOWEST_PRECEDENCE, contextualHeaders);
+  }
+
+  public CallbackWebGraphQLInterceptor(
+      SubscriptionCallbackHandler subscriptionCallbackHandler,
+      int order,
+      Set<String> contextualHeaders) {
     this.subscriptionCallbackHandler = subscriptionCallbackHandler;
     this.order = order;
+    this.contextualHeaders = contextualHeaders;
   }
 
   @Override
@@ -58,8 +79,10 @@ public Mono<WebGraphQlResponse> intercept(WebGraphQlRequest request, Chain chain
                 if (logger.isDebugEnabled()) {
                   logger.debug("Starting subscription using callback: " + callback);
                 }
+                var callbackWithContextualData =
+                    callback.withContext(parseContextualHeaders(request));
                 return this.subscriptionCallbackHandler
-                    .handleSubscriptionUsingCallback(request, callback)
+                    .handleSubscriptionUsingCallback(request, callbackWithContextualData)
                     .map(response -> callbackResponse(request, response));
               })
           .onErrorResume(
@@ -104,4 +127,20 @@ private WebGraphQlResponse errorCallbackResponse(WebGraphQlRequest request) {
   public int getOrder() {
     return order;
   }
+
+  private Map<String, List<String>> parseContextualHeaders(WebGraphQlRequest request) {
+    if (!this.contextualHeaders.isEmpty()) {
+      Map<String, List<String>> contextualHeaders = new HashMap<>();
+      var headers = request.getHeaders();
+      for (String header : this.contextualHeaders) {
+        if (headers.containsKey(header)) {
+          var headerValue = request.getHeaders().get(header);
+          contextualHeaders.put(header, headerValue);
+        }
+      }
+      return contextualHeaders;
+    } else {
+      return new HashMap<>();
+    }
+  }
 }

spring-subscription-callback/src/main/java/com/apollographql/subscription/callback/SubscriptionCallback.java

@@ -2,6 +2,8 @@
 
 import com.apollographql.subscription.exception.CallbackExtensionNotSpecifiedException;
 import com.apollographql.subscription.exception.InvalidCallbackExtensionException;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import org.jetbrains.annotations.NotNull;
 import reactor.core.publisher.Mono;
@@ -13,12 +15,29 @@
  * @param subscription_id The generated unique ID for the subscription operation
  * @param verifier A string that Emitter will include in all HTTP callback requests to verify its
  *     identity
+ * @param heartbeatIntervalMs Interval between heartbeats in milliseconds
+ * @param context Contextual data that should be returned with callbacks as HTTP headers
  */
 public record SubscriptionCallback(
     @NotNull String callback_url,
     @NotNull String subscription_id,
     @NotNull String verifier,
-    int heartbeatIntervalMs) {
+    int heartbeatIntervalMs,
+    @NotNull Map<String, List<String>> context) {
+
+  public SubscriptionCallback(
+      @NotNull String callback_url,
+      @NotNull String subscription_id,
+      @NotNull String verifier,
+      int heartbeatIntervalMs) {
+    this(callback_url, subscription_id, verifier, heartbeatIntervalMs, new HashMap<>());
+  }
+
+  @NotNull
+  public SubscriptionCallback withContext(@NotNull Map<String, List<String>> context) {
+    return new SubscriptionCallback(
+        callback_url, subscription_id, verifier, heartbeatIntervalMs, context);
+  }
 
   public static String SUBSCRIPTION_EXTENSION = "subscription";
   public static String CALLBACK_URL = "callbackUrl";

spring-subscription-callback/src/main/java/com/apollographql/subscription/callback/SubscriptionCallbackHandler.java

@@ -59,6 +59,7 @@ public Mono<ExecutionResult> handleSubscriptionUsingCallback(
         .post()
         .header("Content-Type", MediaType.APPLICATION_JSON_VALUE)
         .header(SUBSCRIPTION_PROTOCOL_HEADER, SUBSCRIPTION_PROTOCOL_HEADER_VALUE)
+        .headers(httpHeaders -> httpHeaders.putAll(callback.context()))
         .bodyValue(checkMessage)
         .exchangeToMono(
             checkResponse -> {
@@ -149,6 +150,7 @@ protected Flux<SubscritionCallbackMessage> startSubscription(
                         .post()
                         .header("Content-Type", MediaType.APPLICATION_JSON_VALUE)
                         .header(SUBSCRIPTION_PROTOCOL_HEADER, SUBSCRIPTION_PROTOCOL_HEADER_VALUE)
+                        .headers(httpHeaders -> httpHeaders.putAll(callback.context()))
                         .bodyValue(message)
                         .exchangeToMono(
                             (routerResponse) -> {
@@ -184,6 +186,7 @@ private Flux<SubscritionCallbackMessage> heartbeatFlux(
                     .post()
                     .header("Content-Type", MediaType.APPLICATION_JSON_VALUE)
                     .header(SUBSCRIPTION_PROTOCOL_HEADER, SUBSCRIPTION_PROTOCOL_HEADER_VALUE)
+                    .headers(httpHeaders -> httpHeaders.putAll(callback.context()))
                     .bodyValue(heartbeat)
                     .exchangeToFlux(
                         (heartBeatResponse) -> {

spring-subscription-callback/src/test/java/com/apollographql/subscription/CallbackGraphQLHttpHandlerAbstrIT.java

@@ -4,11 +4,15 @@
 import static com.apollographql.subscription.callback.SubscriptionCallbackHandler.SUBSCRIPTION_PROTOCOL_HEADER;
 import static com.apollographql.subscription.callback.SubscriptionCallbackHandler.SUBSCRIPTION_PROTOCOL_HEADER_VALUE;
 
+import com.jayway.jsonpath.JsonPath;
 import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.TimeUnit;
 import okhttp3.mockwebserver.MockResponse;
 import okhttp3.mockwebserver.MockWebServer;
+import org.junit.Assert;
 import org.springframework.graphql.test.tester.WebSocketGraphQlTester;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
@@ -18,6 +22,8 @@
 import reactor.test.StepVerifier;
 
 public abstract class CallbackGraphQLHttpHandlerAbstrIT {
+  public static String TEST_HEADER_NAME = "X-Test-Header";
+  public static String TEST_HEADER_VALUE = "junitTEST123";
 
   public void verifyQueriesWorks(WebTestClient testClient) {
     var graphQLRequest = Map.of("query", "query { hello }");
@@ -59,13 +65,23 @@ public void verifyPostSubscriptionsWithoutCallbackDontWork(WebTestClient testCli
         .is4xxClientError();
   }
 
-  public void verifySuccessfulCallbackInit(WebTestClient testClient) {
+  public void verifySuccessfulCallbackSubscription(WebTestClient testClient) {
+    verifyCallbackSubscription(testClient, false);
+  }
+
+  public void verifySuccessfulCallbackSubscriptionWithHeaders(WebTestClient testClient) {
+    verifyCallbackSubscription(testClient, true);
+  }
+
+  private void verifyCallbackSubscription(WebTestClient testClient, boolean testHeaders) {
     try (MockWebServer server = new MockWebServer()) {
       server.enqueue(
           new MockResponse()
               .setResponseCode(HttpStatus.NO_CONTENT.value())
               .setHeader(SUBSCRIPTION_PROTOCOL_HEADER, SUBSCRIPTION_PROTOCOL_HEADER_VALUE));
       server.enqueue(new MockResponse().setResponseCode(HttpStatus.NO_CONTENT.value()));
+      server.enqueue(new MockResponse().setResponseCode(HttpStatus.NO_CONTENT.value()));
+      server.enqueue(new MockResponse().setResponseCode(HttpStatus.NO_CONTENT.value()));
       server.enqueue(new MockResponse().setResponseCode(HttpStatus.NOT_FOUND.value()));
 
       // Start the server.
@@ -79,14 +95,49 @@ public void verifySuccessfulCallbackInit(WebTestClient testClient) {
           .post()
           .uri("/graphql")
           .contentType(MediaType.APPLICATION_JSON)
+          .headers(
+              headers -> {
+                if (testHeaders) {
+                  headers.put(TEST_HEADER_NAME, List.of(TEST_HEADER_VALUE));
+                }
+              })
           .bodyValue(graphQLRequest)
           .exchange()
           .expectStatus()
           .is2xxSuccessful()
           .expectHeader()
           .valueEquals(SUBSCRIPTION_PROTOCOL_HEADER, SUBSCRIPTION_PROTOCOL_HEADER_VALUE);
+
+      var checkMessageRaw = server.takeRequest(1, TimeUnit.SECONDS);
+      Assert.assertNotNull(checkMessageRaw);
+      if (testHeaders) {
+        var headerValue = checkMessageRaw.getHeader(TEST_HEADER_NAME);
+        Assert.assertEquals(TEST_HEADER_VALUE, headerValue);
+      }
+      var checkMessage = JsonPath.parse(checkMessageRaw.getBody().inputStream());
+      Assert.assertEquals(subscriptionId, checkMessage.read("$.id"));
+
+      for (int i : List.of(1, 2, 3)) {
+        var nextMessageRaw = server.takeRequest(1, TimeUnit.SECONDS);
+        Assert.assertNotNull(nextMessageRaw);
+        if (testHeaders) {
+          var headerValue = checkMessageRaw.getHeader(TEST_HEADER_NAME);
+          Assert.assertEquals(TEST_HEADER_VALUE, headerValue);
+        }
+        var nextMessage = JsonPath.parse(nextMessageRaw.getBody().inputStream());
+        Assert.assertEquals(subscriptionId, nextMessage.read("$.id"));
+        Assert.assertEquals(Integer.valueOf(i), nextMessage.read("$.payload.data.counter"));
+      }
+
+      var completeMessageRaw = server.takeRequest(1, TimeUnit.SECONDS);
+      Assert.assertNotNull(completeMessageRaw);
+      var completeMessage = JsonPath.parse(completeMessageRaw.getBody().inputStream());
+      Assert.assertEquals(subscriptionId, completeMessage.read("$.id"));
+      Assert.assertEquals("complete", completeMessage.read("$.action"));
     } catch (IOException e) {
       // failed to close the server
+    } catch (InterruptedException e) {
+      Assert.fail("Failed to read callback data. Exception: " + e.getMessage());
     }
   }
 

spring-subscription-callback/src/test/java/com/apollographql/subscription/configuration/TestGraphQLConfiguration.java

@@ -1,31 +1,18 @@
 package com.apollographql.subscription.configuration;
 
-import com.apollographql.federation.graphqljava.Federation;
-import com.apollographql.federation.graphqljava._Entity;
-import graphql.schema.DataFetcher;
-import java.util.List;
-import java.util.Map;
-import java.util.stream.Collectors;
 import org.springframework.boot.autoconfigure.graphql.GraphQlSourceBuilderCustomizer;
 import org.springframework.context.annotation.Bean;
-import org.springframework.graphql.execution.ClassNameTypeResolver;
+import org.springframework.graphql.data.federation.FederationSchemaFactory;
 
 public class TestGraphQLConfiguration {
-  // configure federation
+
   @Bean
-  public GraphQlSourceBuilderCustomizer federationTransform() {
-    DataFetcher<?> entityDataFetcher =
-        env -> {
-          List<Map<String, Object>> representations = env.getArgument(_Entity.argumentName);
-          return representations.stream().map(representation -> null).collect(Collectors.toList());
-        };
+  public GraphQlSourceBuilderCustomizer customizer(FederationSchemaFactory factory) {
+    return builder -> builder.schemaFactory(factory::createGraphQLSchema);
+  }
 
-    return builder ->
-        builder.schemaFactory(
-            (registry, wiring) ->
-                Federation.transform(registry, wiring)
-                    .fetchEntities(entityDataFetcher)
-                    .resolveEntityType(new ClassNameTypeResolver())
-                    .build());
+  @Bean
+  public FederationSchemaFactory federationSchemaFactory() {
+    return new FederationSchemaFactory();
   }
 }

spring-subscription-callback/src/test/java/com/apollographql/subscription/webflux/CallbackGraphQLWebfluxHttpHandlerIT.java

@@ -4,6 +4,7 @@
 import com.apollographql.subscription.CallbackWebGraphQLInterceptor;
 import com.apollographql.subscription.callback.SubscriptionCallbackHandler;
 import com.apollographql.subscription.configuration.TestGraphQLConfiguration;
+import java.util.Set;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -39,7 +40,8 @@ public SubscriptionCallbackHandler callbackHandler(ExecutionGraphQlService graph
     @Bean
     public CallbackWebGraphQLInterceptor callbackGraphQlInterceptor(
         SubscriptionCallbackHandler callbackHandler) {
-      return new CallbackWebGraphQLInterceptor(callbackHandler);
+      var headers = Set.of(TEST_HEADER_NAME);
+      return new CallbackWebGraphQLInterceptor(callbackHandler, headers);
     }
   }
 
@@ -58,13 +60,18 @@ public void webSocketSubscription_works() {
   }
 
   @Test
-  public void postSubscription_withoutCallback_returns404() {
-    verifyPostSubscriptionsWithoutCallbackDontWork(testClient);
+  public void callbackSubscription_works() {
+    verifySuccessfulCallbackSubscription(testClient);
   }
 
   @Test
-  public void callback_initSuccessful_returns200() {
-    verifySuccessfulCallbackInit(testClient);
+  public void callbackSubscription_withHeaders_works() {
+    verifySuccessfulCallbackSubscriptionWithHeaders(testClient);
+  }
+
+  @Test
+  public void postSubscription_withoutCallback_returns404() {
+    verifyPostSubscriptionsWithoutCallbackDontWork(testClient);
   }
 
   @Test