Add onRawStatus to WebClient

- Add onRawStatus to WebClient.ResponseSpec, allowing users to deal with raw status codes that are not in HttpStatus. - No longer throw an exception status codes not in HttpStatus. Closes gh-23367
spring-projects · Jul 30, 2019 · 6cb4b8b · 6cb4b8b
1 parent 7b69726
commit 6cb4b8b
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 27 deletions.
diff --git a/...flux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java b/...flux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java
@@ -28,6 +28,7 @@
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.IntPredicate;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
@@ -396,8 +397,10 @@ public HttpHeaders getHeaders() {
 
 	private static class DefaultResponseSpec implements ResponseSpec {
 
+		private static final IntPredicate STATUS_CODE_ERROR = value -> value >= 400;
+
 		private static final StatusHandler DEFAULT_STATUS_HANDLER =
-				new StatusHandler(HttpStatus::isError, DefaultResponseSpec::createResponseException);
+				new StatusHandler(STATUS_CODE_ERROR, DefaultResponseSpec::createResponseException);
 
 		private final Mono<ClientResponse> responseMono;
 
@@ -414,11 +417,24 @@ private static class DefaultResponseSpec implements ResponseSpec {
 		@Override
 		public ResponseSpec onStatus(Predicate<HttpStatus> statusPredicate,
 				Function<ClientResponse, Mono<? extends Throwable>> exceptionFunction) {
+			return onRawStatus(toIntPredicate(statusPredicate), exceptionFunction);
+		}
+
+		private static IntPredicate toIntPredicate(Predicate<HttpStatus> predicate) {
+			return value -> {
+				HttpStatus status = HttpStatus.resolve(value);
+				return (status != null) && predicate.test(status);
+			};
+		}
+
+		@Override
+		public ResponseSpec onRawStatus(IntPredicate statusCodePredicate,
+				Function<ClientResponse, Mono<? extends Throwable>> exceptionFunction) {
 
 			if (this.statusHandlers.size() == 1 && this.statusHandlers.get(0) == DEFAULT_STATUS_HANDLER) {
 				this.statusHandlers.clear();
 			}
-			this.statusHandlers.add(new StatusHandler(statusPredicate,
+			this.statusHandlers.add(new StatusHandler(statusCodePredicate,
 					(clientResponse, request) -> exceptionFunction.apply(clientResponse)));
 
 			return this;
@@ -451,27 +467,23 @@ public <T> Flux<T> bodyToFlux(ParameterizedTypeReference<T> elementType) {
 		private <T extends Publisher<?>> T handleBody(ClientResponse response,
 				T bodyPublisher, Function<Mono<? extends Throwable>, T> errorFunction) {
 
-			if (HttpStatus.resolve(response.rawStatusCode()) != null) {
-				for (StatusHandler handler : this.statusHandlers) {
-					if (handler.test(response.statusCode())) {
-						HttpRequest request = this.requestSupplier.get();
-						Mono<? extends Throwable> exMono;
-						try {
-							exMono = handler.apply(response, request);
-							exMono = exMono.flatMap(ex -> drainBody(response, ex));
-							exMono = exMono.onErrorResume(ex -> drainBody(response, ex));
-						}
-						catch (Throwable ex2) {
-							exMono = drainBody(response, ex2);
-						}
-						return errorFunction.apply(exMono);
+			int statusCode = response.rawStatusCode();
+			for (StatusHandler handler : this.statusHandlers) {
+				if (handler.test(statusCode)) {
+					HttpRequest request = this.requestSupplier.get();
+					Mono<? extends Throwable> exMono;
+					try {
+						exMono = handler.apply(response, request);
+						exMono = exMono.flatMap(ex -> drainBody(response, ex));
+						exMono = exMono.onErrorResume(ex -> drainBody(response, ex));
 					}
+					catch (Throwable ex2) {
+						exMono = drainBody(response, ex2);
+					}
+					return errorFunction.apply(exMono);
 				}
-				return bodyPublisher;
-			}
-			else {
-				return errorFunction.apply(createResponseException(response, this.requestSupplier.get()));
 			}
+			return bodyPublisher;
 		}
 
 		@SuppressWarnings("unchecked")
@@ -520,11 +532,11 @@ private static Mono<WebClientResponseException> createResponseException(
 
 		private static class StatusHandler {
 
-			private final Predicate<HttpStatus> predicate;
+			private final IntPredicate predicate;
 
 			private final BiFunction<ClientResponse, HttpRequest, Mono<? extends Throwable>> exceptionFunction;
 
-			public StatusHandler(Predicate<HttpStatus> predicate,
+			public StatusHandler(IntPredicate predicate,
 					BiFunction<ClientResponse, HttpRequest, Mono<? extends Throwable>> exceptionFunction) {
 
 				Assert.notNull(predicate, "Predicate must not be null");
@@ -533,13 +545,15 @@ public StatusHandler(Predicate<HttpStatus> predicate,
 				this.exceptionFunction = exceptionFunction;
 			}
 
-			public boolean test(HttpStatus status) {
+			public boolean test(int status) {
 				return this.predicate.test(status);
 			}
 
 			public Mono<? extends Throwable> apply(ClientResponse response, HttpRequest request) {
 				return this.exceptionFunction.apply(response, request);
 			}
+
+
 		}
 	}
 

diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2019 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.
@@ -23,6 +23,7 @@
 import java.util.Map;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.IntPredicate;
 import java.util.function.Predicate;
 
 import org.reactivestreams.Publisher;
@@ -591,7 +592,7 @@ interface ResponseSpec {
 		 * Register a custom error function that gets invoked when the given {@link HttpStatus}
 		 * predicate applies. The exception returned from the function will be returned from
 		 * {@link #bodyToMono(Class)} and {@link #bodyToFlux(Class)}.
-		 * <p>By default, an error handler is register that throws a
+		 * <p>By default, an error handler is registered that throws a
 		 * {@link WebClientResponseException} when the response status code is 4xx or 5xx.
 		 * @param statusPredicate a predicate that indicates whether {@code exceptionFunction}
 		 * applies
@@ -604,6 +605,24 @@ interface ResponseSpec {
 		ResponseSpec onStatus(Predicate<HttpStatus> statusPredicate,
 				Function<ClientResponse, Mono<? extends Throwable>> exceptionFunction);
 
+		/**
+		 * Register a custom error function that gets invoked when the given raw status code
+		 * predicate applies. The exception returned from the function will be returned from
+		 * {@link #bodyToMono(Class)} and {@link #bodyToFlux(Class)}.
+		 * <p>By default, an error handler is registered that throws a
+		 * {@link WebClientResponseException} when the response status code is 4xx or 5xx.
+		 * @param statusCodePredicate a predicate of the raw status code that indicates
+		 * whether {@code exceptionFunction} applies.
+		 * <p><strong>NOTE:</strong> if the response is expected to have content,
+		 * the exceptionFunction should consume it. If not, the content will be
+		 * automatically drained to ensure resources are released.
+		 * @param exceptionFunction the function that returns the exception
+		 * @return this builder
+		 * @since 5.1.9
+		 */
+		ResponseSpec onRawStatus(IntPredicate statusCodePredicate,
+				Function<ClientResponse, Mono<? extends Throwable>> exceptionFunction);
+
 		/**
 		 * Extract the body to a {@code Mono}. By default, if the response has status code 4xx or
 		 * 5xx, the {@code Mono} will contain a {@link WebClientException}. This can be overridden

diff --git a/...test/java/org/springframework/web/reactive/function/client/WebClientIntegrationTests.java b/...test/java/org/springframework/web/reactive/function/client/WebClientIntegrationTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2019 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.
@@ -56,7 +56,11 @@
 import org.springframework.http.client.reactive.ReactorClientHttpConnector;
 import org.springframework.http.codec.Pojo;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Integration tests using an {@link ExchangeFunction} through {@link WebClient}.
@@ -606,6 +610,28 @@ public void shouldApplyCustomStatusHandler() {
 		});
 	}
 
+	@Test
+	public void shouldApplyCustomRawStatusHandler() {
+		prepareResponse(response -> response.setResponseCode(500)
+				.setHeader("Content-Type", "text/plain").setBody("Internal Server error"));
+
+		Mono<String> result = this.webClient.get()
+				.uri("/greeting?name=Spring")
+				.retrieve()
+				.onRawStatus(value -> value >= 500 && value < 600, response -> Mono.just(new MyException("500 error!")))
+				.bodyToMono(String.class);
+
+		StepVerifier.create(result)
+				.expectError(MyException.class)
+				.verify(Duration.ofSeconds(3));
+
+		expectRequestCount(1);
+		expectRequest(request -> {
+			assertEquals("*/*", request.getHeader(HttpHeaders.ACCEPT));
+			assertEquals("/greeting?name=Spring", request.getPath());
+		});
+	}
+
 	@Test
 	public void shouldApplyCustomStatusHandlerParameterizedTypeReference() {
 		prepareResponse(response -> response.setResponseCode(500)