Autoconfigure reactive subscriptions for spring webflux

Author: oliemansm

example-request-scoped-dataloader/src/main/java/graphql/servlet/examples/dataloader/requestscope/CustomerResolver.java

@@ -1,21 +1,21 @@
 package graphql.servlet.examples.dataloader.requestscope;
 
 import com.coxautodev.graphql.tools.GraphQLResolver;
+import graphql.kickstart.execution.context.GraphQLContext;
 import graphql.schema.DataFetchingEnvironment;
-import graphql.servlet.context.GraphQLContext;
+import java.util.concurrent.CompletableFuture;
 import org.dataloader.DataLoader;
 import org.springframework.stereotype.Component;
 
-import java.util.concurrent.CompletableFuture;
-
 @Component
 public class CustomerResolver implements GraphQLResolver<Customer> {
 
-    public CompletableFuture<String> getName(Customer customer, DataFetchingEnvironment dfe) {
-        final DataLoader<Integer, String> dataloader = ((GraphQLContext) dfe.getContext())
-                .getDataLoaderRegistry().get()
-                .getDataLoader("customerDataLoader");
+  public CompletableFuture<String> getName(Customer customer, DataFetchingEnvironment dfe) {
+    final DataLoader<Integer, String> dataloader = ((GraphQLContext) dfe.getContext())
+        .getDataLoaderRegistry().get()
+        .getDataLoader("customerDataLoader");
+
+    return dataloader.load(customer.getCustomerId());
+  }
 
-        return dataloader.load(customer.getCustomerId());
-    }
 }

graphql-kickstart-spring-boot-autoconfigure-webflux/src/main/java/graphql/kickstart/spring/webflux/boot/GraphQLSpringWebfluxAutoConfiguration.java

@@ -7,19 +7,36 @@
 import graphql.kickstart.execution.GraphQLObjectMapper;
 import graphql.kickstart.execution.config.GraphQLBuilder;
 import graphql.kickstart.execution.config.ObjectMapperProvider;
-import graphql.kickstart.spring.DefaultGraphQLSpringInvocationInputFactory;
-import graphql.kickstart.spring.GraphQLSpringContextBuilder;
-import graphql.kickstart.spring.GraphQLSpringInvocationInputFactory;
-import graphql.kickstart.spring.GraphQLSpringRootObjectBuilder;
+import graphql.kickstart.execution.subscriptions.GraphQLSubscriptionInvocationInputFactory;
+import graphql.kickstart.execution.subscriptions.apollo.ApolloSubscriptionConnectionListener;
+import graphql.kickstart.execution.subscriptions.apollo.KeepAliveSubscriptionConnectionListener;
+import graphql.kickstart.spring.webflux.DefaultGraphQLSpringWebfluxContextBuilder;
+import graphql.kickstart.spring.webflux.DefaultGraphQLSpringWebfluxRootObjectBuilder;
 import graphql.kickstart.spring.webflux.GraphQLController;
+import graphql.kickstart.spring.webflux.GraphQLSpringWebfluxContextBuilder;
+import graphql.kickstart.spring.webflux.GraphQLSpringWebfluxInvocationInputFactory;
+import graphql.kickstart.spring.webflux.GraphQLSpringWebfluxRootObjectBuilder;
+import graphql.kickstart.spring.webflux.ReactiveSubscriptionsProtocolFactory;
+import graphql.kickstart.spring.webflux.ReactiveWebSocketSubscriptionsHandler;
+import graphql.kickstart.spring.webflux.apollo.ReactiveApolloSubscriptionProtocolFactory;
 import graphql.schema.GraphQLSchema;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.ComponentScan;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.web.reactive.HandlerMapping;
+import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
+import org.springframework.web.reactive.socket.WebSocketHandler;
+import org.springframework.web.reactive.socket.server.support.WebSocketHandlerAdapter;
 
 @Slf4j
 @Configuration
@@ -37,11 +54,23 @@ public GraphQLObjectMapper graphQLObjectMapper(ObjectProvider<ObjectMapperProvid
 
   @Bean
   @ConditionalOnMissingBean
-  public GraphQLSpringInvocationInputFactory graphQLSpringInvocationInputFactory(
-      @Autowired(required = false) GraphQLSpringContextBuilder contextBuilder,
-      @Autowired(required = false) GraphQLSpringRootObjectBuilder rootObjectBuilder
+  public GraphQLSpringWebfluxContextBuilder graphQLSpringWebfluxContextBuilder() {
+    return new DefaultGraphQLSpringWebfluxContextBuilder();
+  }
+
+  @Bean
+  @ConditionalOnMissingBean
+  public GraphQLSpringWebfluxRootObjectBuilder graphQLSpringWebfluxRootObjectBuilder() {
+    return new DefaultGraphQLSpringWebfluxRootObjectBuilder();
+  }
+
+  @Bean
+  @ConditionalOnMissingBean
+  public GraphQLSpringWebfluxInvocationInputFactory graphQLSpringInvocationInputFactory(
+      @Autowired(required = false) GraphQLSpringWebfluxContextBuilder contextBuilder,
+      @Autowired(required = false) GraphQLSpringWebfluxRootObjectBuilder rootObjectBuilder
   ) {
-    return new DefaultGraphQLSpringInvocationInputFactory(contextBuilder, rootObjectBuilder);
+    return new GraphQLSpringWebfluxInvocationInputFactory(contextBuilder, rootObjectBuilder);
   }
 
   @Bean
@@ -51,4 +80,46 @@ public GraphQLInvoker graphQLInvoker(GraphQLSchema schema) {
     return new GraphQLInvoker(graphQL);
   }
 
+  @Bean
+  @ConditionalOnMissingBean
+  public ReactiveSubscriptionsProtocolFactory subscriptionProtocolFactory(
+      GraphQLObjectMapper graphQLObjectMapper,
+      GraphQLSubscriptionInvocationInputFactory invocationInputFactory,
+      GraphQLInvoker graphQLInvoker,
+      @Autowired(required = false) Collection<ApolloSubscriptionConnectionListener> connectionListeners
+  ) {
+    Set<ApolloSubscriptionConnectionListener> listeners = new HashSet<>();
+    if (connectionListeners != null) {
+      listeners.addAll(connectionListeners);
+    }
+    if (listeners.stream().noneMatch(KeepAliveSubscriptionConnectionListener.class::isInstance)) {
+      listeners.add(new KeepAliveSubscriptionConnectionListener());
+    }
+    return new ReactiveApolloSubscriptionProtocolFactory(
+        graphQLObjectMapper,
+        invocationInputFactory,
+        graphQLInvoker,
+        listeners
+    );
+  }
+
+  @Bean
+  public HandlerMapping webSocketHandlerMapping(
+      @Value("${graphql.subscriptions.url:subscriptions}") String path,
+      ReactiveWebSocketSubscriptionsHandler webSocketHandler) {
+    Map<String, WebSocketHandler> map = new HashMap<>();
+    map.put(path, webSocketHandler);
+
+    SimpleUrlHandlerMapping handlerMapping = new SimpleUrlHandlerMapping();
+    handlerMapping.setOrder(1);
+    handlerMapping.setUrlMap(map);
+    return handlerMapping;
+  }
+
+  @Bean
+  @ConditionalOnMissingBean
+  WebSocketHandlerAdapter webSocketHandlerAdapter() {
+    return new WebSocketHandlerAdapter();
+  }
+
 }

graphql-kickstart-spring-support/src/main/java/graphql/kickstart/spring/DefaultGraphQLSpringInvocationInputFactory.java

@@ -9,7 +9,7 @@
 
 public class DefaultGraphQLSpringInvocationInputFactory implements GraphQLSpringInvocationInputFactory {
 
-  private Supplier<GraphQLSpringContextBuilder> contextBuilderSupplier = () -> (DefaultGraphQLSpringContext::new);
+  private Supplier<GraphQLSpringContextBuilder> contextBuilderSupplier = () -> (GraphQLSpringServerWebExchangeContext::new);
   private Supplier<GraphQLSpringRootObjectBuilder> rootObjectBuilderSupplier = () -> (serverWebExchange -> new Object());
 
   public DefaultGraphQLSpringInvocationInputFactory(
@@ -52,4 +52,12 @@ public GraphQLBatchedInvocationInput create(Collection<GraphQLRequest> graphQLRe
     throw new UnsupportedOperationException("Batch queries not suppoprted yet");
   }
 
+  protected Supplier<GraphQLSpringContextBuilder> getContextBuilderSupplier() {
+    return contextBuilderSupplier;
+  }
+
+  protected Supplier<GraphQLSpringRootObjectBuilder> getRootObjectBuilderSupplier() {
+    return rootObjectBuilderSupplier;
+  }
+
 }

graphql-kickstart-spring-support/src/main/java/graphql/kickstart/spring/DefaultGraphQLSpringContext.java renamed to graphql-kickstart-spring-support/src/main/java/graphql/kickstart/spring/GraphQLSpringServerWebExchangeContext.java

@@ -5,15 +5,15 @@
 import org.dataloader.DataLoaderRegistry;
 import org.springframework.web.server.ServerWebExchange;
 
-public class DefaultGraphQLSpringContext extends DefaultGraphQLContext implements GraphQLSpringContext {
+public class GraphQLSpringServerWebExchangeContext extends DefaultGraphQLContext implements GraphQLSpringContext {
 
   private final ServerWebExchange serverWebExchange;
 
-  public DefaultGraphQLSpringContext(ServerWebExchange serverWebExchange) {
+  public GraphQLSpringServerWebExchangeContext(ServerWebExchange serverWebExchange) {
     this(new DataLoaderRegistry(), serverWebExchange);
   }
 
-  public DefaultGraphQLSpringContext(DataLoaderRegistry dataLoaderRegistry, ServerWebExchange serverWebExchange) {
+  public GraphQLSpringServerWebExchangeContext(DataLoaderRegistry dataLoaderRegistry, ServerWebExchange serverWebExchange) {
     super(dataLoaderRegistry, null);
     this.serverWebExchange = Objects.requireNonNull(serverWebExchange, "Server web exchange cannot be null");
   }

graphql-kickstart-spring-webflux/src/main/java/graphql/kickstart/spring/webflux/DefaultGraphQLSpringWebSocketSessionContext.java

@@ -0,0 +1,28 @@
+package graphql.kickstart.spring.webflux;
+
+import graphql.kickstart.execution.context.DefaultGraphQLContext;
+import java.util.Objects;
+import org.dataloader.DataLoaderRegistry;
+import org.springframework.web.reactive.socket.WebSocketSession;
+
+public class DefaultGraphQLSpringWebSocketSessionContext extends DefaultGraphQLContext implements
+    GraphQLSpringWebSocketSessionContext {
+
+  private final WebSocketSession webSocketSession;
+
+  public DefaultGraphQLSpringWebSocketSessionContext(WebSocketSession webSocketSession) {
+    this(new DataLoaderRegistry(), webSocketSession);
+  }
+
+  public DefaultGraphQLSpringWebSocketSessionContext(DataLoaderRegistry dataLoaderRegistry,
+      WebSocketSession webSocketSession) {
+    super(dataLoaderRegistry, null);
+    this.webSocketSession = Objects.requireNonNull(webSocketSession, "WebSocketSession is required");
+  }
+
+  @Override
+  public WebSocketSession getWebSocketSession() {
+    return webSocketSession;
+  }
+
+}

graphql-kickstart-spring-webflux/src/main/java/graphql/kickstart/spring/webflux/DefaultGraphQLSpringWebfluxContextBuilder.java

@@ -0,0 +1,20 @@
+package graphql.kickstart.spring.webflux;
+
+import graphql.kickstart.spring.GraphQLSpringContext;
+import graphql.kickstart.spring.GraphQLSpringServerWebExchangeContext;
+import org.springframework.web.reactive.socket.WebSocketSession;
+import org.springframework.web.server.ServerWebExchange;
+
+public class DefaultGraphQLSpringWebfluxContextBuilder implements GraphQLSpringWebfluxContextBuilder {
+
+  @Override
+  public GraphQLSpringWebSocketSessionContext build(WebSocketSession webSocketSession) {
+    return new DefaultGraphQLSpringWebSocketSessionContext(webSocketSession);
+  }
+
+  @Override
+  public GraphQLSpringContext build(ServerWebExchange serverWebExchange) {
+    return new GraphQLSpringServerWebExchangeContext(serverWebExchange);
+  }
+
+}

graphql-kickstart-spring-webflux/src/main/java/graphql/kickstart/spring/webflux/DefaultGraphQLSpringWebfluxRootObjectBuilder.java

@@ -0,0 +1,20 @@
+package graphql.kickstart.spring.webflux;
+
+import graphql.kickstart.spring.GraphQLSpringRootObjectBuilder;
+import org.springframework.web.reactive.socket.WebSocketSession;
+import org.springframework.web.server.ServerWebExchange;
+
+public class DefaultGraphQLSpringWebfluxRootObjectBuilder
+    implements GraphQLSpringRootObjectBuilder, GraphQLSpringWebfluxRootObjectBuilder {
+
+  @Override
+  public Object build(ServerWebExchange serverWebExchange) {
+    return new Object();
+  }
+
+  @Override
+  public Object build(WebSocketSession webSocketSession) {
+    return new Object();
+  }
+
+}

graphql-kickstart-spring-webflux/src/main/java/graphql/kickstart/spring/webflux/GraphQLSpringWebSocketSessionContext.java

@@ -0,0 +1,10 @@
+package graphql.kickstart.spring.webflux;
+
+import graphql.kickstart.execution.context.GraphQLContext;
+import org.springframework.web.reactive.socket.WebSocketSession;
+
+interface GraphQLSpringWebSocketSessionContext extends GraphQLContext {
+
+  WebSocketSession getWebSocketSession();
+
+}

graphql-kickstart-spring-webflux/src/main/java/graphql/kickstart/spring/webflux/GraphQLSpringWebfluxContextBuilder.java

@@ -0,0 +1,10 @@
+package graphql.kickstart.spring.webflux;
+
+import graphql.kickstart.spring.GraphQLSpringContextBuilder;
+import org.springframework.web.reactive.socket.WebSocketSession;
+
+public interface GraphQLSpringWebfluxContextBuilder extends GraphQLSpringContextBuilder {
+
+  GraphQLSpringWebSocketSessionContext build(WebSocketSession webSocketSession);
+
+}

graphql-kickstart-spring-webflux/src/main/java/graphql/kickstart/spring/webflux/GraphQLSpringWebfluxInvocationInputFactory.java

@@ -0,0 +1,38 @@
+package graphql.kickstart.spring.webflux;
+
+import graphql.kickstart.execution.GraphQLRequest;
+import graphql.kickstart.execution.input.GraphQLSingleInvocationInput;
+import graphql.kickstart.execution.subscriptions.GraphQLSubscriptionInvocationInputFactory;
+import graphql.kickstart.execution.subscriptions.SubscriptionSession;
+import graphql.kickstart.spring.DefaultGraphQLSpringInvocationInputFactory;
+import graphql.kickstart.spring.GraphQLSpringContextBuilder;
+import graphql.kickstart.spring.GraphQLSpringRootObjectBuilder;
+import java.util.function.Supplier;
+import org.springframework.web.reactive.socket.WebSocketSession;
+
+public class GraphQLSpringWebfluxInvocationInputFactory
+    extends DefaultGraphQLSpringInvocationInputFactory
+    implements GraphQLSubscriptionInvocationInputFactory {
+
+  public GraphQLSpringWebfluxInvocationInputFactory(GraphQLSpringContextBuilder contextBuilder,
+      GraphQLSpringRootObjectBuilder rootObjectBuilder) {
+    super(contextBuilder, rootObjectBuilder);
+  }
+
+  public GraphQLSpringWebfluxInvocationInputFactory(
+      Supplier<GraphQLSpringContextBuilder> contextBuilderSupplier,
+      Supplier<GraphQLSpringRootObjectBuilder> rootObjectBuilderSupplier) {
+    super(contextBuilderSupplier, rootObjectBuilderSupplier);
+  }
+
+  @Override
+  public GraphQLSingleInvocationInput create(GraphQLRequest graphQLRequest, SubscriptionSession session) {
+    return new GraphQLSingleInvocationInput(
+        graphQLRequest,
+        null,
+        ((GraphQLSpringWebfluxContextBuilder) getContextBuilderSupplier().get()).build((WebSocketSession) session.unwrap()),
+        ((GraphQLSpringWebfluxRootObjectBuilder) getRootObjectBuilderSupplier().get()).build((WebSocketSession) session.unwrap())
+    );
+  }
+
+}