DATAES-719 - Add customization hook for reactive WebClient.

Original PR: #363

src/main/asciidoc/reference/elasticsearch-clients.adoc

@@ -108,8 +108,15 @@ static class Config {
   @Bean
   ReactiveElasticsearchClient client() {
 
-    ClientConfiguration clientConfiguration = ClientConfiguration.builder() <1>
+    ClientConfiguration clientConfiguration = ClientConfiguration.builder()   <1>
       .connectedTo("localhost:9200", "localhost:9291")
+      .withWebClientConfigurer(webClient -> {                                 <2>
+        ExchangeStrategies exchangeStrategies = ExchangeStrategies.builder()
+            .codecs(configurer -> configurer.defaultCodecs()
+                .maxInMemorySize(-1))
+            .build();
+        return webClient.mutate().exchangeStrategies(exchangeStrategies).build();
+       })
       .build();
 
     return ReactiveRestClients.create(clientConfiguration);
@@ -128,6 +135,7 @@ Mono<IndexResponse> response = client.index(request ->
 );
 ----
 <1> Use the builder to provide cluster addresses, set default `HttpHeaders` or enable SSL.
+<2> when configuring a reactive client, the `withWebClientConfigurer` hook can be used to customize the WebClient.
 ====
 
 NOTE: The ReactiveClient response, especially for search operations, is bound to the `from` (offset) & `size` (limit) options of the request.

src/main/java/org/springframework/data/elasticsearch/client/ClientConfiguration.java

@@ -20,11 +20,13 @@ import java.net.SocketAddress;
 import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Function;
 
 import javax.net.ssl.HostnameVerifier;
 import javax.net.ssl.SSLContext;
 
 import org.springframework.http.HttpHeaders;
+import org.springframework.web.reactive.function.client.WebClient;
 
 /**
  * Configuration interface exposing common client configuration properties for Elasticsearch clients.
@@ -161,6 +163,11 @@ public interface ClientConfiguration {
 	 */
 	Optional<String> getProxy();
 
+	/**
+	 * @return the function for configuring a WebClient.
+	 */
+	Function<WebClient, WebClient> getWebClientConfigurer();
+
 	/**
 	 * @author Christoph Strobl
 	 */
@@ -314,9 +321,17 @@ public interface ClientConfiguration {
 
 		/**
 		 * @param proxy a proxy formatted as String {@literal host:port}.
-		 * @return the {@link MaybeSecureClientConfigurationBuilder}.
+		 * @return the {@link TerminalClientConfigurationBuilder}.
 		 */
-		MaybeSecureClientConfigurationBuilder withProxy(String proxy);
+		TerminalClientConfigurationBuilder withProxy(String proxy);
+
+		/**
+		 * set customization hook in case of a reactive configuration
+		 * 
+		 * @param webClientConfigurer function to configure the WebClient
+		 * @return the {@link TerminalClientConfigurationBuilder}.
+		 */
+		TerminalClientConfigurationBuilder withWebClientConfigurer(Function<WebClient, WebClient> webClientConfigurer);
 
 		/**
 		 * Build the {@link ClientConfiguration} object.

src/main/java/org/springframework/data/elasticsearch/client/ClientConfigurationBuilder.java

@@ -20,6 +20,7 @@ import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import javax.net.ssl.HostnameVerifier;
@@ -31,6 +32,7 @@ import org.springframework.data.elasticsearch.client.ClientConfiguration.Termina
 import org.springframework.http.HttpHeaders;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.web.reactive.function.client.WebClient;
 
 /**
  * Default builder implementation for {@link ClientConfiguration}.
@@ -56,6 +58,7 @@ class ClientConfigurationBuilder
 	private String password;
 	private String pathPrefix;
 	private String proxy;
+	private Function<WebClient, WebClient> webClientConfigurer;
 
 	/*
 	 * (non-Javadoc)
@@ -187,12 +190,20 @@ class ClientConfigurationBuilder
 
 	@Override
 	public TerminalClientConfigurationBuilder withPathPrefix(String pathPrefix) {
-
 		this.pathPrefix = pathPrefix;
 
 		return this;
 	}
 
+	@Override
+	public TerminalClientConfigurationBuilder withWebClientConfigurer(Function<WebClient, WebClient> webClientConfigurer) {
+
+		Assert.notNull(webClientConfigurer, "webClientConfigurer must not be null");
+
+		this.webClientConfigurer = webClientConfigurer;
+		return this;
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration.ClientConfigurationBuilderWithOptionalDefaultHeaders#build()
@@ -208,7 +219,7 @@ class ClientConfigurationBuilder
 		}
 
 		return new DefaultClientConfiguration(hosts, headers, useSsl, sslContext, soTimeout, connectTimeout, pathPrefix,
-				hostnameVerifier, proxy);
+				hostnameVerifier, proxy, webClientConfigurer);
 	}
 
 	private static InetSocketAddress parse(String hostAndPort) {

src/main/java/org/springframework/data/elasticsearch/client/DefaultClientConfiguration.java

@@ -21,12 +21,14 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Function;
 
 import javax.net.ssl.HostnameVerifier;
 import javax.net.ssl.SSLContext;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.lang.Nullable;
+import org.springframework.web.reactive.function.client.WebClient;
 
 /**
  * Default {@link ClientConfiguration} implementation.
@@ -34,6 +36,7 @@ import org.springframework.lang.Nullable;
  * @author Mark Paluch
  * @author Christoph Strobl
  * @author Huw Ayling-Miller
+ * @author Peter-Josef Meisch
  * @since 3.2
  */
 class DefaultClientConfiguration implements ClientConfiguration {
@@ -47,10 +50,11 @@ class DefaultClientConfiguration implements ClientConfiguration {
 	private final String pathPrefix;
 	private final @Nullable HostnameVerifier hostnameVerifier;
 	private final String proxy;
+	private final Function<WebClient, WebClient> webClientConfigurer;
 
 	DefaultClientConfiguration(List<InetSocketAddress> hosts, HttpHeaders headers, boolean useSsl,
 			@Nullable SSLContext sslContext, Duration soTimeout, Duration connectTimeout, @Nullable String pathPrefix,
-			@Nullable HostnameVerifier hostnameVerifier, String proxy) {
+			@Nullable HostnameVerifier hostnameVerifier, String proxy, Function<WebClient, WebClient> webClientConfigurer) {
 
 		this.hosts = Collections.unmodifiableList(new ArrayList<>(hosts));
 		this.headers = new HttpHeaders(headers);
@@ -61,86 +65,56 @@ class DefaultClientConfiguration implements ClientConfiguration {
 		this.pathPrefix = pathPrefix;
 		this.hostnameVerifier = hostnameVerifier;
 		this.proxy = proxy;
+		this.webClientConfigurer = webClientConfigurer;
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration#getEndpoints()
-	 */
 	@Override
 	public List<InetSocketAddress> getEndpoints() {
 		return this.hosts;
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration#getDefaultHeaders()
-	 */
 	@Override
 	public HttpHeaders getDefaultHeaders() {
 		return this.headers;
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration#useSsl()
-	 */
 	@Override
 	public boolean useSsl() {
 		return this.useSsl;
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration#getSslContext()
-	 */
 	@Override
 	public Optional<SSLContext> getSslContext() {
 		return Optional.ofNullable(this.sslContext);
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration#getHostNameVerifier()
-	 */
 	@Override
 	public Optional<HostnameVerifier> getHostNameVerifier() {
 		return Optional.ofNullable(this.hostnameVerifier);
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration#getConnectTimeout()
-	 */
 	@Override
 	public Duration getConnectTimeout() {
 		return this.connectTimeout;
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration#getSocketTimeout()
-	 */
 	@Override
 	public Duration getSocketTimeout() {
 		return this.soTimeout;
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration#getPathPrefix()
-	 */
 	@Override
 	public String getPathPrefix() {
 		return this.pathPrefix;
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration#getProxy()
-	 */
 	@Override
 	public Optional<String> getProxy() {
 		return Optional.ofNullable(proxy);
 	}
+
+	@Override
+	public Function<WebClient, WebClient> getWebClientConfigurer() {
+		return webClientConfigurer != null ? webClientConfigurer : Function.identity();
+	}
 }

src/main/java/org/springframework/data/elasticsearch/client/reactive/DefaultReactiveElasticsearchClient.java

@@ -226,7 +226,9 @@ public class DefaultReactiveElasticsearchClient implements ReactiveElasticsearch
 			provider = provider.withPathPrefix(clientConfiguration.getPathPrefix());
 		}
 
-		return provider.withDefaultHeaders(clientConfiguration.getDefaultHeaders());
+		provider = provider.withDefaultHeaders(clientConfiguration.getDefaultHeaders()) //
+				.withWebClientConfigurer(clientConfiguration.getWebClientConfigurer());
+		return provider;
 	}
 
 	/*

src/main/java/org/springframework/data/elasticsearch/client/reactive/DefaultWebClientProvider.java

@@ -19,6 +19,7 @@ import java.net.InetSocketAddress;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.client.reactive.ClientHttpConnector;
@@ -33,6 +34,7 @@ import org.springframework.web.reactive.function.client.WebClient.Builder;
  * @author Mark Paluch
  * @author Christoph Strobl
  * @author Huw Ayling-Miller
+ * @author Peter-Josef Meisch
  * @since 3.2
  */
 class DefaultWebClientProvider implements WebClientProvider {
@@ -44,6 +46,7 @@ class DefaultWebClientProvider implements WebClientProvider {
 	private final Consumer<Throwable> errorListener;
 	private final HttpHeaders headers;
 	private final String pathPrefix;
+	private final Function<WebClient, WebClient> webClientConfigurer;
 
 	/**
 	 * Create new {@link DefaultWebClientProvider} with empty {@link HttpHeaders} and no-op {@literal error listener}.
@@ -52,24 +55,28 @@ class DefaultWebClientProvider implements WebClientProvider {
 	 * @param connector can be {@literal null}.
 	 */
 	DefaultWebClientProvider(String scheme, @Nullable ClientHttpConnector connector) {
-		this(scheme, connector, e -> {}, HttpHeaders.EMPTY, null);
+		this(scheme, connector, e -> {}, HttpHeaders.EMPTY, null, Function.identity());
 	}
 
 	/**
 	 * Create new {@link DefaultWebClientProvider} with empty {@link HttpHeaders} and no-op {@literal error listener}.
 	 *
+	 * @param pathPrefixcan be {@literal null}
 	 * @param scheme must not be {@literal null}.
 	 * @param connector can be {@literal null}.
 	 * @param errorListener must not be {@literal null}.
 	 * @param headers must not be {@literal null}.
-     * @param pathPrefixcan be {@literal null}
+	 * @param webClientConfigurer must not be {@literal null}.
 	 */
 	private DefaultWebClientProvider(String scheme, @Nullable ClientHttpConnector connector,
-			Consumer<Throwable> errorListener, HttpHeaders headers, @Nullable String pathPrefix) {
+			Consumer<Throwable> errorListener, HttpHeaders headers, @Nullable String pathPrefix,
+			Function<WebClient, WebClient> webClientConfigurer) {
 
 		Assert.notNull(scheme, "Scheme must not be null! A common scheme would be 'http'.");
-		Assert.notNull(errorListener, "ErrorListener must not be null! You may want use a no-op one 'e -> {}' instead.");
+		Assert.notNull(errorListener, "errorListener must not be null! You may want use a no-op one 'e -> {}' instead.");
 		Assert.notNull(headers, "headers must not be null! Think about using 'HttpHeaders.EMPTY' as an alternative.");
+		Assert.notNull(webClientConfigurer,
+				"webClientConfigurer must not be null! You may want use a no-op one 'Function.identity()' instead.");
 
 		this.cachedClients = new ConcurrentHashMap<>();
 		this.scheme = scheme;
@@ -77,12 +84,9 @@ class DefaultWebClientProvider implements WebClientProvider {
 		this.errorListener = errorListener;
 		this.headers = headers;
 		this.pathPrefix = pathPrefix;
+		this.webClientConfigurer = webClientConfigurer;
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.reactive.WebClientProvider#get(java.net.InetSocketAddress)
-	 */
 	@Override
 	public WebClient get(InetSocketAddress endpoint) {
 
@@ -91,19 +95,21 @@ class DefaultWebClientProvider implements WebClientProvider {
 		return this.cachedClients.computeIfAbsent(endpoint, this::createWebClientForSocketAddress);
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.reactive.WebClientProvider#getDefaultHeaders()
-	 */
 	@Override
 	public HttpHeaders getDefaultHeaders() {
 		return headers;
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.reactive.WebClientProvider#withDefaultHeaders(org.springframework.http.HttpHeaders)
-	 */
+	@Override
+	public Consumer<Throwable> getErrorListener() {
+		return this.errorListener;
+	}
+
+	@Override
+	public String getPathPrefix() {
+		return pathPrefix;
+	}
+
 	@Override
 	public WebClientProvider withDefaultHeaders(HttpHeaders headers) {
 
@@ -113,51 +119,33 @@ class DefaultWebClientProvider implements WebClientProvider {
 		merged.addAll(this.headers);
 		merged.addAll(headers);
 
-		return new DefaultWebClientProvider(this.scheme, this.connector, errorListener, merged, this.pathPrefix);
+		return new DefaultWebClientProvider(scheme, connector, errorListener, merged, pathPrefix, webClientConfigurer);
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.reactive.WebClientProvider#getErrorListener()
-	 */
-	@Override
-	public Consumer<Throwable> getErrorListener() {
-		return this.errorListener;
-	}
-
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.reactive.WebClientProvider#getPathPrefix()
-	 */
-	@Override
-	public String getPathPrefix() {
-		return pathPrefix;
-	}
-
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.reactive.WebClientProvider#withErrorListener(java.util.function.Consumer)
-	 */
 	@Override
 	public WebClientProvider withErrorListener(Consumer<Throwable> errorListener) {
 
 		Assert.notNull(errorListener, "Error listener must not be null.");
 
 		Consumer<Throwable> listener = this.errorListener.andThen(errorListener);
-		return new DefaultWebClientProvider(this.scheme, this.connector, listener, this.headers, this.pathPrefix);
+		return new DefaultWebClientProvider(scheme, this.connector, listener, headers, pathPrefix, webClientConfigurer);
 	}
 
-	/*
-	 * (non-Javadoc)
-	 * @see org.springframework.data.elasticsearch.client.reactive.WebClientProvider#withPathPrefix(java.lang.String)
-	 */
 	@Override
 	public WebClientProvider withPathPrefix(String pathPrefix) {
 		Assert.notNull(pathPrefix, "pathPrefix must not be null.");
 
-		return new DefaultWebClientProvider(this.scheme, this.connector, this.errorListener, this.headers, pathPrefix);
+		return new DefaultWebClientProvider(this.scheme, this.connector, this.errorListener, this.headers, pathPrefix,
+				webClientConfigurer);
 	}
 
+	@Override
+	public WebClientProvider withWebClientConfigurer(Function<WebClient, WebClient> webClientConfigurer) {
+		return new DefaultWebClientProvider(scheme, connector, errorListener, headers, pathPrefix, webClientConfigurer);
+
+	}
+
+
 	protected WebClient createWebClientForSocketAddress(InetSocketAddress socketAddress) {
 
 		Builder builder = WebClient.builder().defaultHeaders(it -> it.addAll(getDefaultHeaders()));
@@ -167,7 +155,8 @@ class DefaultWebClientProvider implements WebClientProvider {
 		}
 
 		String baseUrl = String.format("%s://%s:%d%s", this.scheme, socketAddress.getHostString(), socketAddress.getPort(),
-				pathPrefix == null ? "" : "/" + pathPrefix);
-		return builder.baseUrl(baseUrl).filter((request, next) -> next.exchange(request).doOnError(errorListener)).build();
+				pathPrefix == null ? "" : '/' + pathPrefix);
+		WebClient webClient = builder.baseUrl(baseUrl).filter((request, next) -> next.exchange(request).doOnError(errorListener)).build();
+		return webClientConfigurer.apply(webClient);
 	}
 }

src/main/java/org/springframework/data/elasticsearch/client/reactive/WebClientProvider.java

@@ -17,6 +17,7 @@ package org.springframework.data.elasticsearch.client.reactive;
 
 import java.net.InetSocketAddress;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.client.reactive.ClientHttpConnector;
@@ -35,6 +36,7 @@ import org.springframework.web.reactive.function.client.WebClient;
  * @author Christoph Strobl
  * @author Mark Paluch
  * @author Huw Ayling-Miller
+ * @author Peter-Josef Meisch
  * @since 3.2
  */
 public interface WebClientProvider {
@@ -129,4 +131,12 @@ public interface WebClientProvider {
 	 * @since 4.0
 	 */
 	WebClientProvider withPathPrefix(String pathPrefix);
+
+	/**
+	 * Create a new instance of {@link WebClientProvider} calling the given {@link Function} to configure the {@link WebClient}.
+	 * @param webClientConfigurer configuration function
+	 * @return new instance of {@link WebClientProvider}
+	 * @since 4.0
+	 */
+	WebClientProvider withWebClientConfigurer(Function<WebClient, WebClient> webClientConfigurer);
 }

src/test/java/org/springframework/data/elasticsearch/client/ClientConfigurationUnitTests.java

@@ -20,12 +20,14 @@ import static org.mockito.Mockito.*;
 
 import java.net.InetSocketAddress;
 import java.time.Duration;
+import java.util.function.Function;
 
 import javax.net.ssl.SSLContext;
 
 import org.apache.http.conn.ssl.NoopHostnameVerifier;
 import org.junit.jupiter.api.Test;
 import org.springframework.http.HttpHeaders;
+import org.springframework.web.reactive.function.client.WebClient;
 
 /**
  * Unit tests for {@link ClientConfiguration}.
@@ -143,6 +145,26 @@ public class ClientConfigurationUnitTests {
 		assertThat(clientConfiguration.getHostNameVerifier()).contains(NoopHostnameVerifier.INSTANCE);
 	}
 
+	@Test // DATAES-719
+	void shouldHaveDefaultWebClientConfigurer() {
+		ClientConfiguration clientConfiguration = ClientConfiguration.builder() //
+				.connectedTo("foo", "bar") //
+				.build();
+
+		assertThat(clientConfiguration.getWebClientConfigurer()).isEqualTo(Function.identity());
+	}
+
+	@Test // DATAES-719
+	void shouldUseConfiguredWebClientConfigurer() {
+		Function<WebClient, WebClient> webClientConfigurer = webClient -> webClient;
+		ClientConfiguration clientConfiguration = ClientConfiguration.builder() //
+				.connectedTo("foo", "bar") //
+				.withWebClientConfigurer(webClientConfigurer) //
+				.build();
+
+		assertThat(clientConfiguration.getWebClientConfigurer()).isEqualTo(webClientConfigurer);
+	}
+
 	private static String buildBasicAuth(String username, String password) {
 
 		HttpHeaders headers = new HttpHeaders();

src/test/java/org/springframework/data/elasticsearch/client/reactive/DefaultWebClientProviderUnitTests.java

@@ -18,12 +18,15 @@ package org.springframework.data.elasticsearch.client.reactive;
 import static org.assertj.core.api.Assertions.*;
 
 import java.net.InetSocketAddress;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
 
 import org.junit.jupiter.api.Test;
 import org.springframework.web.reactive.function.client.WebClient;
 
 /**
  * @author Christoph Strobl
+ * @author Peter-Josef Meisch
  */
 public class DefaultWebClientProviderUnitTests {
 
@@ -40,4 +43,18 @@ public class DefaultWebClientProviderUnitTests {
 		assertThat(shouldBeCachedInstanceOfClient1).isSameAs(client1);
 		assertThat(notClient1ButAnotherInstance).isNotSameAs(client1);
 	}
+
+	@Test // DATAES-719
+	void shouldCallWebClientConfigurer() {
+		AtomicReference<Boolean> configurerCalled = new AtomicReference<>(false);
+		Function<WebClient, WebClient> configurer = webClient -> {
+			configurerCalled.set(true);
+			return webClient;
+		};
+		WebClientProvider provider = new DefaultWebClientProvider("http", null).withWebClientConfigurer(configurer);
+
+		provider.get(InetSocketAddress.createUnresolved("localhost", 9200));
+
+		assertThat(configurerCalled).hasValue(true);
+	}
 }

src/test/java/org/springframework/data/elasticsearch/client/reactive/ReactiveMockClientTestsUtils.java

@@ -61,6 +61,7 @@ import org.springframework.web.util.UriBuilder;
  * @author Christoph Strobl
  * @author Huw Ayling-Miller
  * @author Henrique Amaral
+ * @author Peter-Josef Meisch
  */
 public class ReactiveMockClientTestsUtils {
 
@@ -276,6 +277,11 @@ public class ReactiveMockClientTestsUtils {
 			throw new UnsupportedOperationException();
 		}
 
+		@Override
+		public WebClientProvider withWebClientConfigurer(Function<WebClient, WebClient> webClientConfigurer) {
+			throw new UnsupportedOperationException("not implemented");
+		}
+
 		public Send when(String host) {
 			InetSocketAddress inetSocketAddress = getInetSocketAddress(host);
 			return new CallbackImpl(get(host), headersUriSpecMap.get(inetSocketAddress),