Add WebSocketGraphQlClientInterceptor

See gh-322

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/client/AbstractGraphQlClientBuilder.java

@@ -111,6 +111,14 @@ protected void setJsonCodecs(Encoder<?> encoder, Decoder<?> decoder) {
 		this.jsonDecoder = decoder;
 	}
 
+	/**
+	 * Return the configured interceptors. For subclasses that look for a
+	 * transport specific interceptor extensions.
+	 */
+	protected List<GraphQlClientInterceptor> getInterceptors() {
+		return this.interceptors;
+	}
+
 	/**
 	 * Build the default transport-agnostic client that subclasses can then wrap
 	 * with {@link AbstractDelegatingGraphQlClient}.

spring-graphql/src/main/java/org/springframework/graphql/client/DefaultWebSocketGraphQlClient.java

@@ -18,7 +18,9 @@
 
 import java.net.URI;
 import java.util.Arrays;
+import java.util.List;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 import reactor.core.publisher.Mono;
 
@@ -157,12 +159,25 @@ public WebSocketGraphQlClient build() {
 					CodecDelegate.findJsonDecoder(this.codecConfigurer));
 
 			WebSocketGraphQlTransport transport = new WebSocketGraphQlTransport(
-					this.url, this.headers, this.webSocketClient, this.codecConfigurer, null, payload -> {});
+					this.url, this.headers, this.webSocketClient, this.codecConfigurer, getInterceptor());
 
 			GraphQlClient graphQlClient = super.buildGraphQlClient(transport);
 			return new DefaultWebSocketGraphQlClient(graphQlClient, transport, getBuilderInitializer());
 		}
 
+		private WebSocketGraphQlClientInterceptor getInterceptor() {
+
+			List<WebSocketGraphQlClientInterceptor> interceptors = getInterceptors().stream()
+					.filter(interceptor -> interceptor instanceof WebSocketGraphQlClientInterceptor)
+					.map(interceptor -> (WebSocketGraphQlClientInterceptor) interceptor)
+					.collect(Collectors.toList());
+
+			Assert.state(interceptors.size() <= 1,
+					"Only a single interceptor of type WebSocketGraphQlClientInterceptor may be configured");
+
+			return (!interceptors.isEmpty() ? interceptors.get(0) : new WebSocketGraphQlClientInterceptor() {});
+		}
+
 	}
 
 }

spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketGraphQlClientInterceptor.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2020-2022 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.client;
+
+
+import java.util.Map;
+
+import reactor.core.publisher.Mono;
+
+
+/**
+ * An extension of {@link GraphQlClientInterceptor} with additional methods to
+ * for WebSocket interception points. Only a single interceptor of type
+ * {@link WebSocketGraphQlClientInterceptor} may be configured.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.0.0
+ */
+public interface WebSocketGraphQlClientInterceptor extends GraphQlClientInterceptor {
+
+	/**
+	 * Provide a {@code Mono} that returns the payload for the
+	 * {@code "connection_init"} message. The {@code Mono} is subscribed to every
+	 * type a new WebSocket connection is established.
+	 */
+	default Mono<Object> connectionInitPayload() {
+		return Mono.empty();
+	}
+
+	/**
+	 * Handler the {@code "connection_ack"} message received from the server at
+	 * the start of the WebSocket connection.
+	 * @param ackPayload the payload of the {@code "connection_ack"} message
+	 */
+	default Mono<Void> handleConnectionAck(Map<String, Object> ackPayload) {
+		return Mono.empty();
+	}
+
+}

spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketGraphQlTransport.java

@@ -22,7 +22,6 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.function.Consumer;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -69,17 +68,18 @@ final class WebSocketGraphQlTransport implements GraphQlTransport {
 
 	WebSocketGraphQlTransport(
 			URI url, @Nullable HttpHeaders headers, WebSocketClient client, CodecConfigurer codecConfigurer,
-			@Nullable Object connectionInitPayload, Consumer<Map<String, Object>> connectionAckHandler) {
+			WebSocketGraphQlClientInterceptor interceptor) {
 
 		Assert.notNull(url, "URI is required");
-		Assert.notNull(url, "URI is required");
+		Assert.notNull(client, "WebSocketClient is required");
+		Assert.notNull(codecConfigurer, "CodecConfigurer is required");
+		Assert.notNull(interceptor, "WebSocketGraphQlClientInterceptor is required");
 
 		this.url = url;
 		this.headers.putAll(headers != null ? headers : HttpHeaders.EMPTY);
 		this.webSocketClient = client;
 
-		this.graphQlSessionHandler = new GraphQlSessionHandler(
-				codecConfigurer, connectionInitPayload, connectionAckHandler);
+		this.graphQlSessionHandler = new GraphQlSessionHandler(codecConfigurer, interceptor);
 
 		this.graphQlSessionMono = initGraphQlSession(this.url, this.headers, client, this.graphQlSessionHandler)
 				.cacheInvalidateWhen(GraphQlSession::notifyWhenClosed);
@@ -167,21 +167,16 @@ private static class GraphQlSessionHandler implements WebSocketHandler {
 
 		private final CodecDelegate codecDelegate;
 
-		private final GraphQlMessage connectionInitMessage;
-
-		private final Consumer<Map<String, Object>> connectionAckHandler;
+		private final WebSocketGraphQlClientInterceptor interceptor;
 
 		private Sinks.One<GraphQlSession> graphQlSessionSink;
 
 		private final AtomicBoolean stopped = new AtomicBoolean();
 
 
-		GraphQlSessionHandler(CodecConfigurer codecConfigurer,
-				@Nullable Object connectionInitPayload, Consumer<Map<String, Object>> connectionAckHandler) {
-
+		GraphQlSessionHandler(CodecConfigurer codecConfigurer, WebSocketGraphQlClientInterceptor interceptor) {
 			this.codecDelegate = new CodecDelegate(codecConfigurer);
-			this.connectionInitMessage = GraphQlMessage.connectionInit(connectionInitPayload);
-			this.connectionAckHandler = connectionAckHandler;
+			this.interceptor = interceptor;
 			this.graphQlSessionSink = Sinks.unsafe().one();
 		}
 
@@ -231,8 +226,12 @@ public Mono<Void> handle(WebSocketSession session) {
 			GraphQlSession graphQlSession = new GraphQlSession(session);
 			registerCloseStatusHandling(graphQlSession, session);
 
+			Mono<GraphQlMessage> connectionInitMono = this.interceptor.connectionInitPayload()
+					.defaultIfEmpty(Collections.emptyMap())
+					.map(GraphQlMessage::connectionInit);
+
 			Mono<Void> sendCompletion =
-					session.send(Flux.just(this.connectionInitMessage).concatWith(graphQlSession.getRequestFlux())
+					session.send(connectionInitMono.concatWith(graphQlSession.getRequestFlux())
 							.map(message -> this.codecDelegate.encode(session, message)));
 
 			Mono<Void> receiveCompletion = session.receive()
@@ -242,20 +241,23 @@ public Mono<Void> handle(WebSocketSession session) {
 								GraphQlMessage message = this.codecDelegate.decode(webSocketMessage);
 								Assert.state(message.resolvedType() == GraphQlMessageType.CONNECTION_ACK,
 										() -> "Unexpected message before connection_ack: " + message);
-								this.connectionAckHandler.accept(message.getPayload());
-								if (logger.isDebugEnabled()) {
-									logger.debug(graphQlSession + " initialized");
-								}
+								return this.interceptor.handleConnectionAck(message.getPayload())
+										.then(Mono.defer(() -> {
+											if (logger.isDebugEnabled()) {
+												logger.debug(graphQlSession + " initialized");
+											}
+											Sinks.EmitResult result = this.graphQlSessionSink.tryEmitValue(graphQlSession);
+											if (result.isFailure()) {
+												return Mono.error(new IllegalStateException(
+														"GraphQlSession initialized but could not be emitted: " + result));
+											}
+											return Mono.empty();
+										}));
 							}
 							catch (Throwable ex) {
 								this.graphQlSessionSink.tryEmitError(ex);
 								return Mono.error(ex);
 							}
-							Sinks.EmitResult emitResult = this.graphQlSessionSink.tryEmitValue(graphQlSession);
-							if (emitResult.isFailure()) {
-								return Mono.error(new IllegalStateException(
-										"GraphQlSession initialized but could not be emitted: " + emitResult));
-							}
 						}
 						else {
 							try {

spring-graphql/src/test/java/org/springframework/graphql/client/MockWebSocketGraphQlTransportTests.java

@@ -31,10 +31,10 @@
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
-import org.springframework.graphql.support.DefaultGraphQlRequest;
 import org.springframework.graphql.GraphQlRequest;
 import org.springframework.graphql.GraphQlResponse;
 import org.springframework.graphql.ResponseError;
+import org.springframework.graphql.support.DefaultGraphQlRequest;
 import org.springframework.graphql.web.TestWebSocketClient;
 import org.springframework.graphql.web.TestWebSocketConnection;
 import org.springframework.graphql.web.support.GraphQlMessage;
@@ -196,9 +196,23 @@ void start() {
 		Map<String, String> initPayload = Collections.singletonMap("key", "valueInit");
 		AtomicReference<Map<String, Object>> connectionAckRef = new AtomicReference<>();
 
+		WebSocketGraphQlClientInterceptor interceptor = new WebSocketGraphQlClientInterceptor() {
+
+			@Override
+			public Mono<Object> connectionInitPayload() {
+				return Mono.just(initPayload);
+			}
+
+			@Override
+			public Mono<Void> handleConnectionAck(Map<String, Object> ackPayload) {
+				connectionAckRef.set(ackPayload);
+				return Mono.empty();
+			}
+		};
+
+
 		WebSocketGraphQlTransport transport = new WebSocketGraphQlTransport(
-				URI.create("/"), HttpHeaders.EMPTY, client, ClientCodecConfigurer.create(),
-				initPayload, connectionAckRef::set);
+				URI.create("/"), HttpHeaders.EMPTY, client, ClientCodecConfigurer.create(), interceptor);
 
 		transport.start().block(TIMEOUT);
 
@@ -311,7 +325,8 @@ void errorDuringResponseHandling() {
 
 	private static WebSocketGraphQlTransport createTransport(WebSocketClient client) {
 		return new WebSocketGraphQlTransport(
-				URI.create("/"), HttpHeaders.EMPTY, client, ClientCodecConfigurer.create(), null, p -> {});
+				URI.create("/"), HttpHeaders.EMPTY, client, ClientCodecConfigurer.create(),
+				new WebSocketGraphQlClientInterceptor() {});
 	}
 
 	private void assertActualClientMessages(GraphQlMessage... expectedMessages) {