diff --git a/spring-graphql/src/main/java/org/springframework/graphql/web/DefaultWebGraphQlHandlerBuilder.java b/spring-graphql/src/main/java/org/springframework/graphql/web/DefaultWebGraphQlHandlerBuilder.java index 343c63ff..a5822b60 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/web/DefaultWebGraphQlHandlerBuilder.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/web/DefaultWebGraphQlHandlerBuilder.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2021 the original author or authors. + * Copyright 2002-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. @@ -18,10 +18,7 @@ package org.springframework.graphql.web; import java.util.ArrayList; import java.util.Arrays; -import java.util.Collections; import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; import reactor.core.publisher.Mono; @@ -32,6 +29,7 @@ import org.springframework.lang.Nullable; import org.springframework.util.Assert; import org.springframework.util.CollectionUtils; + /** * Default implementation of {@link WebGraphQlHandler.Builder}. * @@ -41,17 +39,21 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder { private final GraphQlService service; + private final List interceptors = new ArrayList<>(); + @Nullable - private List interceptors; + private WebSocketInterceptor webSocketInterceptor; @Nullable private List accessors; + DefaultWebGraphQlHandlerBuilder(GraphQlService service) { Assert.notNull(service, "GraphQlService is required"); this.service = service; } + @Override public WebGraphQlHandler.Builder interceptor(WebInterceptor... interceptors) { return interceptors(Arrays.asList(interceptors)); @@ -59,10 +61,13 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder { @Override public WebGraphQlHandler.Builder interceptors(List interceptors) { - if (!CollectionUtils.isEmpty(interceptors)) { - this.interceptors = (this.interceptors != null) ? this.interceptors : new ArrayList<>(); - this.interceptors.addAll(interceptors); - } + this.interceptors.addAll(interceptors); + interceptors.forEach(interceptor -> { + if (interceptor instanceof WebSocketInterceptor) { + Assert.isNull(this.webSocketInterceptor, "There can be at most 1 WebSocketInterceptor"); + this.webSocketInterceptor = (WebSocketInterceptor) interceptor; + } + }); return this; } @@ -83,98 +88,34 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder { @Override public WebGraphQlHandler build() { - List interceptorsToUse = - (this.interceptors != null) ? this.interceptors : Collections.emptyList(); + WebInterceptorChain endOfChain = + webInput -> this.service.execute(webInput).map(WebOutput::new); - WebInterceptorChain interceptorChain = initWebInterceptorChain(interceptorsToUse); - WebSocketInterceptor webSocketInterceptor = initWebSocketInterceptor(interceptorsToUse); - - WebGraphQlHandler graphQlHandler = new WebGraphQlHandler() { - - @Override - public Mono handleRequest(WebInput input) { - return interceptorChain.next(input); - } - - @Override - public Mono handleWebSocketInitialization(Map payload) { - return (webSocketInterceptor != null ? - webSocketInterceptor.handleConnectionInitialization(payload) : Mono.empty()); - } - - @Override - public Mono handleWebSocketCompletion() { - return (webSocketInterceptor != null ? - webSocketInterceptor.handleConnectionCompletion() : Mono.empty()); - } - }; - - if (!CollectionUtils.isEmpty(this.accessors)) { - graphQlHandler = new ThreadLocalExtractingHandler( - graphQlHandler, ThreadLocalAccessor.composite(this.accessors)); - } - - return graphQlHandler; - } - - private WebInterceptorChain initWebInterceptorChain(List interceptors) { - - WebInterceptorChain endOfChain = webInput -> this.service.execute(webInput).map(WebOutput::new); - - return interceptors.stream() + WebInterceptorChain chain = this.interceptors.stream() .reduce(WebInterceptor::andThen) .map(interceptor -> (WebInterceptorChain) (input) -> interceptor.intercept(input, endOfChain)) .orElse(endOfChain); - } - @Nullable - private WebSocketInterceptor initWebSocketInterceptor(List interceptors) { + return new WebGraphQlHandler() { - List filtered = interceptors.stream() - .filter(current -> current instanceof WebSocketInterceptor) - .map(current -> (WebSocketInterceptor) current) - .collect(Collectors.toList()); + @Override + public Mono handleRequest(WebInput input) { + return chain.next(input) + .contextWrite(context -> { + if (!CollectionUtils.isEmpty(accessors)) { + ThreadLocalAccessor accessor = ThreadLocalAccessor.composite(accessors); + return ReactorContextManager.extractThreadLocalValues(accessor, context); + } + return context; + }); + } - if (filtered.size() > 1) { - throw new IllegalArgumentException( - "There can be at most 1 WebSocketInterceptor. Found " + filtered.size() + "."); - } - - return (!filtered.isEmpty() ? filtered.get(0) : null); - } - - - /** - * {@link WebGraphQlHandler} that extracts ThreadLocal values and saves them in the - * Reactor context for subsequent use for DataFetcher's. - */ - private static class ThreadLocalExtractingHandler implements WebGraphQlHandler { - - private final WebGraphQlHandler delegate; - - private final ThreadLocalAccessor accessor; - - ThreadLocalExtractingHandler(WebGraphQlHandler delegate, ThreadLocalAccessor accessor) { - this.delegate = delegate; - this.accessor = accessor; - } - - @Override - public Mono handleRequest(WebInput input) { - return this.delegate.handleRequest(input).contextWrite((context) -> - ReactorContextManager.extractThreadLocalValues(this.accessor, context)); - } - - @Override - public Mono handleWebSocketInitialization(Map payload) { - return this.delegate.handleWebSocketInitialization(payload); - } - - @Override - public Mono handleWebSocketCompletion() { - return this.delegate.handleWebSocketCompletion(); - } + @Override + public WebSocketInterceptor webSocketInterceptor() { + return (webSocketInterceptor != null ? webSocketInterceptor : new WebSocketInterceptor() {}); + } + }; } } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/web/WebGraphQlHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/web/WebGraphQlHandler.java index fa889d44..bbe6b4dd 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/web/WebGraphQlHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/web/WebGraphQlHandler.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2021 the original author or authors. + * Copyright 2002-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. @@ -17,13 +17,13 @@ package org.springframework.graphql.web; import java.util.List; -import java.util.Map; import reactor.core.publisher.Mono; import org.springframework.graphql.GraphQlService; import org.springframework.graphql.execution.ThreadLocalAccessor; + /** * Contract for common handling of a GraphQL request over HTTP or WebSocket, * for use with Spring MVC or Spring WebFlux. @@ -33,6 +33,7 @@ import org.springframework.graphql.execution.ThreadLocalAccessor; */ public interface WebGraphQlHandler { + /** * Execute the given request and return the resulting output. * @param input the GraphQL request input container @@ -41,24 +42,10 @@ public interface WebGraphQlHandler { Mono handleRequest(WebInput input); /** - * Handle the payload from the connection initialization message that a - * GraphQL over WebSocket client must send after the WebSocket session is - * established and before sending any requests. - * @param payload the payload from the {@code ConnectionInit} message - * @return an optional payload for the {@code ConnectionAck} message + * Return the single interceptor of type {@link WebSocketInterceptor} among + * all the configured interceptors. */ - default Mono handleWebSocketInitialization(Map payload) { - return Mono.empty(); - } - - /** - * Handle the completion message that a GraphQL over WebSocket clients sends - * before closing the WebSocket connection. - * @return signals the end of completion handling - */ - default Mono handleWebSocketCompletion() { - return Mono.empty(); - } + WebSocketInterceptor webSocketInterceptor(); /** diff --git a/spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandler.java index 0d03a7d4..cd227ea8 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandler.java @@ -35,6 +35,7 @@ import reactor.core.publisher.Mono; import org.springframework.graphql.web.WebGraphQlHandler; import org.springframework.graphql.web.WebInput; import org.springframework.graphql.web.WebOutput; +import org.springframework.graphql.web.WebSocketInterceptor; import org.springframework.http.codec.CodecConfigurer; import org.springframework.util.Assert; import org.springframework.util.CollectionUtils; @@ -61,6 +62,8 @@ public class GraphQlWebSocketHandler implements WebSocketHandler { private final WebGraphQlHandler graphQlHandler; + private final WebSocketInterceptor webSocketInterceptor; + private final CodecDelegate codecDelegate; private final Duration initTimeoutDuration; @@ -77,7 +80,9 @@ public class GraphQlWebSocketHandler implements WebSocketHandler { WebGraphQlHandler graphQlHandler, CodecConfigurer codecConfigurer, Duration connectionInitTimeout) { Assert.notNull(graphQlHandler, "WebGraphQlHandler is required"); + this.graphQlHandler = graphQlHandler; + this.webSocketInterceptor = this.graphQlHandler.webSocketInterceptor(); this.codecDelegate = new CodecDelegate(codecConfigurer); this.initTimeoutDuration = connectionInitTimeout; } @@ -137,12 +142,12 @@ public class GraphQlWebSocketHandler implements WebSocketHandler { subscription.cancel(); } } - return this.graphQlHandler.handleWebSocketCompletion().thenMany(Flux.empty()); + return this.webSocketInterceptor.handleConnectionCompletion().thenMany(Flux.empty()); case "connection_init": if (!connectionInitProcessed.compareAndSet(false, true)) { return GraphQlStatus.close(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS); } - return this.graphQlHandler.handleWebSocketInitialization(payload) + return this.webSocketInterceptor.handleConnectionInitialization(payload) .defaultIfEmpty(Collections.emptyMap()) .map(ackPayload -> this.codecDelegate.encodeConnectionAck(session, ackPayload)) .flux() diff --git a/spring-graphql/src/main/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandler.java index 44effd06..566cffba 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandler.java @@ -45,6 +45,7 @@ import reactor.core.scheduler.Schedulers; import org.springframework.graphql.web.WebGraphQlHandler; import org.springframework.graphql.web.WebInput; import org.springframework.graphql.web.WebOutput; +import org.springframework.graphql.web.WebSocketInterceptor; import org.springframework.graphql.web.webflux.GraphQlWebSocketMessage; import org.springframework.http.HttpHeaders; import org.springframework.http.HttpInputMessage; @@ -78,6 +79,8 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub private final WebGraphQlHandler graphQlHandler; + private final WebSocketInterceptor webSocketInterceptor; + private final Duration initTimeoutDuration; private final HttpMessageConverter converter; @@ -92,12 +95,13 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub * message must be received. */ public GraphQlWebSocketHandler( - WebGraphQlHandler graphQlHandler, HttpMessageConverter converter, - Duration connectionInitTimeout) { + WebGraphQlHandler graphQlHandler, HttpMessageConverter converter, Duration connectionInitTimeout) { Assert.notNull(graphQlHandler, "WebGraphQlHandler is required"); Assert.notNull(converter, "HttpMessageConverter for JSON is required"); + this.graphQlHandler = graphQlHandler; + this.webSocketInterceptor = this.graphQlHandler.webSocketInterceptor(); this.initTimeoutDuration = connectionInitTimeout; this.converter = converter; } @@ -166,14 +170,14 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub subscription.cancel(); } } - this.graphQlHandler.handleWebSocketCompletion().block(Duration.ofSeconds(10)); + this.webSocketInterceptor.handleConnectionCompletion().block(Duration.ofSeconds(10)); return; case "connection_init": if (sessionState.setConnectionInitProcessed()) { GraphQlStatus.closeSession(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS); return; } - this.graphQlHandler.handleWebSocketInitialization(payload) + this.webSocketInterceptor.handleConnectionInitialization(payload) .defaultIfEmpty(Collections.emptyMap()) .publishOn(sessionState.getScheduler()) // Serial blocking send via single thread .doOnNext(ackPayload -> { diff --git a/spring-graphql/src/test/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandlerTests.java index 5844526a..54d0cc3e 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandlerTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandlerTests.java @@ -140,8 +140,8 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport { new WebSocketInterceptor() { @Override - public Mono handleConnectionInitialization(Map payload) { - Object value = payload.get("key"); + public Mono handleConnectionInitialization(Map connectionInitPayload) { + Object value = connectionInitPayload.get("key"); return Mono.just(Collections.singletonMap("key", value + " acknowledged")); } }); @@ -162,7 +162,7 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport { new WebSocketInterceptor() { @Override - public Mono handleConnectionInitialization(Map payload) { + public Mono handleConnectionInitialization(Map connectionInitPayload) { return Mono.error(new IllegalStateException()); } }); diff --git a/spring-graphql/src/test/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandlerTests.java index 1ebeb37c..ecdc9d41 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandlerTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandlerTests.java @@ -141,8 +141,8 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport { WebSocketInterceptor interceptor = new WebSocketInterceptor() { @Override - public Mono handleConnectionInitialization(Map payload) { - Object value = payload.get("key"); + public Mono handleConnectionInitialization(Map connectionInitPayload) { + Object value = connectionInitPayload.get("key"); return Mono.just(Collections.singletonMap("key", value + " acknowledged")); } }; @@ -166,7 +166,7 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport { WebSocketInterceptor interceptor = new WebSocketInterceptor() { @Override - public Mono handleConnectionInitialization(Map payload) { + public Mono handleConnectionInitialization(Map connectionInitPayload) { return Mono.error(new IllegalStateException()); } };