Support for subscriptions with WebMvc and spring-websocket

graphql-spring-boot-starter/src/main/java/org/springframework/graphql/boot/WebMvcGraphQLAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2020 the original author or authors.
+ * Copyright 2020-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
 package org.springframework.graphql.boot;
 
 import java.util.Collections;
+import java.util.Map;
 
 import graphql.GraphQL;
 
@@ -24,15 +25,22 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
+import org.springframework.boot.autoconfigure.http.HttpMessageConverters;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.io.Resource;
 import org.springframework.core.io.ResourceLoader;
 import org.springframework.graphql.webmvc.GraphQLHttpHandler;
+import org.springframework.graphql.webmvc.GraphQLWebSocketHandler;
 import org.springframework.http.MediaType;
+import org.springframework.http.converter.HttpMessageConverter;
+import org.springframework.web.servlet.HandlerMapping;
 import org.springframework.web.servlet.function.RouterFunction;
 import org.springframework.web.servlet.function.RouterFunctions;
 import org.springframework.web.servlet.function.ServerResponse;
+import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
+import org.springframework.web.socket.server.support.DefaultHandshakeHandler;
+import org.springframework.web.socket.server.support.WebSocketHttpRequestHandler;
 
 import static org.springframework.web.servlet.function.RequestPredicates.accept;
 import static org.springframework.web.servlet.function.RequestPredicates.contentType;
@@ -50,6 +58,22 @@ public class WebMvcGraphQLAutoConfiguration {
 		return new GraphQLHttpHandler(graphQLBuilder.build(), Collections.emptyList());
 	}
 
+	@Bean
+	@ConditionalOnMissingBean
+	public GraphQLWebSocketHandler graphQLWebSocketHandler(
+			GraphQL.Builder graphQLBuilder, GraphQLProperties properties, HttpMessageConverters converters) {
+
+		HttpMessageConverter<?> converter = converters.getConverters().stream()
+				.filter(candidate -> candidate.canRead(Map.class, MediaType.APPLICATION_JSON))
+				.findFirst()
+				.orElseThrow(() -> new IllegalStateException("No JSON converter"));
+
+		return new GraphQLWebSocketHandler(
+				graphQLBuilder.build(), Collections.emptyList(),
+				converter, properties.getConnectionInitTimeoutDuration()
+		);
+	}
+
 	@Bean
 	public RouterFunction<ServerResponse> graphQLQueryEndpoint(
 			ResourceLoader resourceLoader, GraphQLHttpHandler handler, GraphQLProperties properties) {
@@ -63,4 +87,16 @@ public class WebMvcGraphQLAutoConfiguration {
 				.build();
 	}
 
+	@Bean
+	public HandlerMapping graphQLWebSocketEndpoint(GraphQLWebSocketHandler handler, GraphQLProperties properties) {
+		WebSocketHttpRequestHandler httpRequestHandler =
+				new WebSocketHttpRequestHandler(handler, new DefaultHandshakeHandler());
+
+		String path = properties.getWebSocketPath();
+		SimpleUrlHandlerMapping mapping = new SimpleUrlHandlerMapping();
+		mapping.setUrlMap(Collections.singletonMap(path, httpRequestHandler));
+		mapping.setOrder(-1); // Ahead of annotated controllers
+		return mapping;
+	}
+
 }