Access to ApplicationContext via ServerWebExchange

Issue: SPR-16298

Author: rstoyanchev

spring-web/src/main/java/org/springframework/web/server/ServerWebExchange.java

@@ -24,6 +24,7 @@
 
 import reactor.core.publisher.Mono;
 
+import org.springframework.context.ApplicationContext;
 import org.springframework.context.i18n.LocaleContext;
 import org.springframework.http.codec.multipart.Part;
 import org.springframework.http.server.reactive.ServerHttpRequest;
@@ -136,6 +137,17 @@ default <T> T getAttributeOrDefault(String name, T defaultValue) {
 	 */
 	LocaleContext getLocaleContext();
 
+	/**
+	 * Return the {@link ApplicationContext} associated with the web application,
+	 * if it was initialized with one via
+	 * {@link org.springframework.web.server.adapter.WebHttpHandlerBuilder#applicationContext
+	 * WebHttpHandlerBuilder#applicationContext}.
+	 * @see org.springframework.web.server.adapter.WebHttpHandlerBuilder#applicationContext(ApplicationContext)
+	 * @since 5.0.3
+	 */
+	@Nullable
+	ApplicationContext getApplicationContext();
+
 	/**
 	 * Returns {@code true} if the one of the {@code checkNotModified} methods
 	 * in this contract were used and they returned true.

spring-web/src/main/java/org/springframework/web/server/ServerWebExchangeDecorator.java

@@ -22,6 +22,7 @@
 
 import reactor.core.publisher.Mono;
 
+import org.springframework.context.ApplicationContext;
 import org.springframework.context.i18n.LocaleContext;
 import org.springframework.http.codec.multipart.Part;
 import org.springframework.http.server.reactive.ServerHttpRequest;
@@ -91,6 +92,11 @@ public LocaleContext getLocaleContext() {
 		return getDelegate().getLocaleContext();
 	}
 
+	@Override
+	public ApplicationContext getApplicationContext() {
+		return getDelegate().getApplicationContext();
+	}
+
 	@Override
 	public Mono<MultiValueMap<String, String>> getFormData() {
 		return getDelegate().getFormData();

spring-web/src/main/java/org/springframework/web/server/adapter/DefaultServerWebExchange.java

@@ -28,6 +28,7 @@
 
 import reactor.core.publisher.Mono;
 
+import org.springframework.context.ApplicationContext;
 import org.springframework.context.i18n.LocaleContext;
 import org.springframework.core.ResolvableType;
 import org.springframework.http.HttpHeaders;
@@ -90,6 +91,9 @@ public class DefaultServerWebExchange implements ServerWebExchange {
 
 	private final Mono<MultiValueMap<String, Part>> multipartDataMono;
 
+	@Nullable
+	private final ApplicationContext applicationContext;
+
 	private volatile boolean notModified;
 
 	private Function<String, String> urlTransformer = url -> url;
@@ -99,6 +103,13 @@ public DefaultServerWebExchange(ServerHttpRequest request, ServerHttpResponse re
 			WebSessionManager sessionManager, ServerCodecConfigurer codecConfigurer,
 			LocaleContextResolver localeContextResolver) {
 
+		this(request, response, sessionManager, codecConfigurer, localeContextResolver, null);
+	}
+
+	DefaultServerWebExchange(ServerHttpRequest request, ServerHttpResponse response,
+			WebSessionManager sessionManager, ServerCodecConfigurer codecConfigurer,
+			LocaleContextResolver localeContextResolver, @Nullable ApplicationContext applicationContext) {
+
 		Assert.notNull(request, "'request' is required");
 		Assert.notNull(response, "'response' is required");
 		Assert.notNull(sessionManager, "'sessionManager' is required");
@@ -111,6 +122,7 @@ public DefaultServerWebExchange(ServerHttpRequest request, ServerHttpResponse re
 		this.localeContextResolver = localeContextResolver;
 		this.formDataMono = initFormData(request, codecConfigurer);
 		this.multipartDataMono = initMultipartData(request, codecConfigurer);
+		this.applicationContext = applicationContext;
 	}
 
 	@SuppressWarnings("unchecked")
@@ -191,11 +203,6 @@ public <T extends Principal> Mono<T> getPrincipal() {
 		return Mono.empty();
 	}
 
-	@Override
-	public LocaleContext getLocaleContext() {
-		return this.localeContextResolver.resolveLocaleContext(this);
-	}
-
 	@Override
 	public Mono<MultiValueMap<String, String>> getFormData() {
 		return this.formDataMono;
@@ -206,6 +213,16 @@ public Mono<MultiValueMap<String, Part>> getMultipartData() {
 		return this.multipartDataMono;
 	}
 
+	@Override
+	public LocaleContext getLocaleContext() {
+		return this.localeContextResolver.resolveLocaleContext(this);
+	}
+
+	@Override
+	public ApplicationContext getApplicationContext() {
+		return this.applicationContext;
+	}
+
 	@Override
 	public boolean isNotModified() {
 		return this.notModified;

spring-web/src/main/java/org/springframework/web/server/adapter/HttpWebHandlerAdapter.java

@@ -24,6 +24,7 @@
 import org.apache.commons.logging.LogFactory;
 import reactor.core.publisher.Mono;
 
+import org.springframework.context.ApplicationContext;
 import org.springframework.core.NestedExceptionUtils;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.codec.ServerCodecConfigurer;
@@ -90,6 +91,9 @@ public class HttpWebHandlerAdapter extends WebHandlerDecorator implements HttpHa
 	@Nullable
 	private LocaleContextResolver localeContextResolver;
 
+	@Nullable
+	private ApplicationContext applicationContext;
+
 
 	public HttpWebHandlerAdapter(WebHandler delegate) {
 		super(delegate);
@@ -126,6 +130,13 @@ public void setCodecConfigurer(ServerCodecConfigurer codecConfigurer) {
 		this.codecConfigurer = codecConfigurer;
 	}
 
+	/**
+	 * Return the configured {@link ServerCodecConfigurer}.
+	 */
+	public ServerCodecConfigurer getCodecConfigurer() {
+		return (this.codecConfigurer != null ? this.codecConfigurer : ServerCodecConfigurer.create());
+	}
+
 	/**
 	 * Configure a custom {@link LocaleContextResolver}. The provided instance is set on
 	 * each created {@link DefaultServerWebExchange}.
@@ -137,13 +148,6 @@ public void setLocaleContextResolver(LocaleContextResolver localeContextResolver
 		this.localeContextResolver = localeContextResolver;
 	}
 
-	/**
-	 * Return the configured {@link ServerCodecConfigurer}.
-	 */
-	public ServerCodecConfigurer getCodecConfigurer() {
-		return (this.codecConfigurer != null ? this.codecConfigurer : ServerCodecConfigurer.create());
-	}
-
 	/**
 	 * Return the configured {@link LocaleContextResolver}.
 	 */
@@ -152,6 +156,27 @@ public LocaleContextResolver getLocaleContextResolver() {
 				this.localeContextResolver : new AcceptHeaderLocaleContextResolver());
 	}
 
+	/**
+	 * Configure the {@code ApplicationContext} associated with the web application,
+	 * if it was initialized with one via
+	 * {@link org.springframework.web.server.adapter.WebHttpHandlerBuilder#applicationContext
+	 * WebHttpHandlerBuilder#applicationContext}.
+	 * @param applicationContext the context
+	 * @since 5.0.3
+	 */
+	public void setApplicationContext(ApplicationContext applicationContext) {
+		this.applicationContext = applicationContext;
+	}
+
+	/**
+	 * Return the configured {@code ApplicationContext}, if any.
+	 * @since 5.0.3
+	 */
+	@Nullable
+	public ApplicationContext getApplicationContext() {
+		return this.applicationContext;
+	}
+
 
 	@Override
 	public Mono<Void> handle(ServerHttpRequest request, ServerHttpResponse response) {
@@ -163,7 +188,7 @@ public Mono<Void> handle(ServerHttpRequest request, ServerHttpResponse response)
 
 	protected ServerWebExchange createExchange(ServerHttpRequest request, ServerHttpResponse response) {
 		return new DefaultServerWebExchange(request, response, this.sessionManager,
-				getCodecConfigurer(), getLocaleContextResolver());
+				getCodecConfigurer(), getLocaleContextResolver(), this.applicationContext);
 	}
 
 	private Mono<Void> handleFailure(ServerHttpResponse response, Throwable ex) {

spring-web/src/main/java/org/springframework/web/server/adapter/WebHttpHandlerBuilder.java

@@ -92,6 +92,9 @@ public class WebHttpHandlerBuilder {
 	@Nullable
 	private LocaleContextResolver localeContextResolver;
 
+	@Nullable
+	private ApplicationContext applicationContext;
+
 
 	/**
 	 * Private constructor.
@@ -101,6 +104,15 @@ private WebHttpHandlerBuilder(WebHandler webHandler) {
 		this.webHandler = webHandler;
 	}
 
+	/**
+	 * Private constructor to use when initialized from an ApplicationContext.
+	 */
+	private WebHttpHandlerBuilder(WebHandler webHandler, ApplicationContext applicationContext) {
+		Assert.notNull(webHandler, "WebHandler must not be null");
+		this.webHandler = webHandler;
+		this.applicationContext = applicationContext;
+	}
+
 	/**
 	 * Copy constructor.
 	 */
@@ -144,8 +156,9 @@ public static WebHttpHandlerBuilder webHandler(WebHandler webHandler) {
 	 * @return the prepared builder
 	 */
 	public static WebHttpHandlerBuilder applicationContext(ApplicationContext context) {
+
 		WebHttpHandlerBuilder builder = new WebHttpHandlerBuilder(
-				context.getBean(WEB_HANDLER_BEAN_NAME, WebHandler.class));
+				context.getBean(WEB_HANDLER_BEAN_NAME, WebHandler.class), context);
 
 		// Autowire lists for @Bean + @Order
 
@@ -275,6 +288,9 @@ public HttpHandler build() {
 		if (this.localeContextResolver != null) {
 			adapted.setLocaleContextResolver(this.localeContextResolver);
 		}
+		if (this.applicationContext != null) {
+			adapted.setApplicationContext(this.applicationContext);
+		}
 
 		return adapted;
 	}

spring-web/src/test/java/org/springframework/web/server/adapter/WebHttpHandlerBuilderTests.java

@@ -38,10 +38,11 @@
 
 import static java.time.Duration.ofMillis;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Unit tests for {@link WebHttpHandlerBuilder}.
- *
  * @author Rossen Stoyanchev
  */
 public class WebHttpHandlerBuilderTests {
@@ -54,6 +55,9 @@ public void orderedWebFilterBeans() throws Exception {
 
 		HttpHandler httpHandler = WebHttpHandlerBuilder.applicationContext(context).build();
 
+		assertTrue(httpHandler instanceof HttpWebHandlerAdapter);
+		assertSame(context, ((HttpWebHandlerAdapter) httpHandler).getApplicationContext());
+
 		MockServerHttpRequest request = MockServerHttpRequest.get("/").build();
 		MockServerHttpResponse response = new MockServerHttpResponse();
 		httpHandler.handle(request, response).block(ofMillis(5000));