DATAES-801 - Implement callback to enable adding custom headers in the REST HTTP request.

Original PR: spring-projects#442

Author: sothawo

pom.xml

@@ -252,7 +252,7 @@
 		<dependency>
 			<groupId>com.github.tomakehurst</groupId>
 			<artifactId>wiremock-jre8</artifactId>
-			<version>2.25.1</version>
+			<version>2.26.3</version>
 			<scope>test</scope>
 			<exclusions>
 				<!-- these exclusions are needed because of Elasticsearch JarHell-->

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

@@ -143,40 +143,47 @@ NOTE: The ReactiveClient response, especially for search operations, is bound to
 [[elasticsearch.clients.configuration]]
 == Client Configuration
 
-Client behaviour can be changed via the `ClientConfiguration` that allows to set options for SSL, connect and socket timeouts.
+Client behaviour can be changed via the `ClientConfiguration` that allows to set options for SSL, connect and socket timeouts, headers and other parameters.
 
 .Client Configuration
 ====
 [source,java]
 ----
-// optional if Basic Auhtentication is needed
 HttpHeaders httpHeaders = new HttpHeaders();
-httpHeaders.add("es-security-runas-user", "some-user")                  <1>
+httpHeaders.add("some-header", "on every request")                      <1>
 
 ClientConfiguration clientConfiguration = ClientConfiguration.builder()
   .connectedTo("localhost:9200", "localhost:9291")                      <2>
-  .withProxy("localhost:8888")                                          <3>
-  .withPathPrefix("ela")                                                <4>
-  .withConnectTimeout(Duration.ofSeconds(5))                            <5>
-  .withSocketTimeout(Duration.ofSeconds(3))                             <6>
-  .useSsl()                                                             <7>
+  .useSsl()                                                             <3>
+  .withProxy("localhost:8888")                                          <4>
+  .withPathPrefix("ela")                                                <5>
+  .withConnectTimeout(Duration.ofSeconds(5))                            <6>
+  .withSocketTimeout(Duration.ofSeconds(3))                             <7>
   .withDefaultHeaders(defaultHeaders)                                   <8>
   .withBasicAuth(username, password)                                    <9>
+  .withHeaders(() -> {                                                  <10>
+    HttpHeaders headers = new HttpHeaders();
+    headers.add("currentTime", LocalDateTime.now().format(DateTimeFormatter.ISO_LOCAL_DATE_TIME));
+    return headers;
+  })
   . // ... other options
   .build();
 
 ----
 <1> Define default headers, if they need to be customized
 <2> Use the builder to provide cluster addresses, set default `HttpHeaders` or enable SSL.
-<3> Optionally set a proxy footnote:notreactive[not yet implemented for the reactive client].
-<4> Optionally set a path prefix, mostly used when different clusters a behind some reverse proxy.
-<5> Set the connection timeout. Default is 10 sec.
-<6> Set the socket timeout. Default is 5 sec.
-<7> Optionally enable SSL.
+<3> Optionally enable SSL.
+<4> Optionally set a proxy.
+<5> Optionally set a path prefix, mostly used when different clusters a behind some reverse proxy.
+<6> Set the connection timeout. Default is 10 sec.
+<7> Set the socket timeout. Default is 5 sec.
 <8> Optionally set headers.
 <9> Add basic authentication.
+<10> A `Supplier<Header>` function can be specified which is called every time before a request is sent to Elasticsearch - here, as an example, the current time is written in a header.
 ====
 
+IMPORTANT: Adding a Header supplier as shown in above example allows to inject headers that may change over the time, like authentication JWT tokens. If this is used in the reactive setup, the supplier function *must not* block!
+
 [[elasticsearch.clients.logging]]
 == Client Logging
 

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

@@ -21,6 +21,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import javax.net.ssl.HostnameVerifier;
 import javax.net.ssl.SSLContext;
@@ -170,6 +171,11 @@ static ClientConfiguration create(InetSocketAddress socketAddress) {
 	 */
 	Function<WebClient, WebClient> getWebClientConfigurer();
 
+	/**
+	 * @return the supplier for custom headers.
+	 */
+	Supplier<HttpHeaders> getHeadersSupplier();
+
 	/**
 	 * @author Christoph Strobl
 	 */
@@ -335,6 +341,19 @@ default TerminalClientConfigurationBuilder withSocketTimeout(long millis) {
 		 */
 		TerminalClientConfigurationBuilder withWebClientConfigurer(Function<WebClient, WebClient> webClientConfigurer);
 
+		/**
+		 * set a supplier for custom headers. This is invoked for every HTTP request to Elasticsearch to retrieve headers
+		 * that should be sent with the request. A common use case is passing in authentication headers that may change.
+		 * <br/>
+		 * Note: When used in a reactive environment, the calling of {@link Supplier#get()} function must not do any
+		 * blocking operations. It may return {@literal null}.
+		 * 
+		 * @param headers supplier function for headers, must not be {@literal null}
+		 * @return the {@link TerminalClientConfigurationBuilder}.
+		 * @since 4.0
+		 */
+		TerminalClientConfigurationBuilder withHeaders(Supplier<HttpHeaders> headers);
+
 		/**
 		 * Build the {@link ClientConfiguration} object.
 		 *

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

@@ -21,6 +21,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 import javax.net.ssl.HostnameVerifier;
@@ -58,7 +59,8 @@ class ClientConfigurationBuilder
 	private @Nullable String password;
 	private @Nullable String pathPrefix;
 	private @Nullable String proxy;
-	private @Nullable Function<WebClient, WebClient> webClientConfigurer;
+	private Function<WebClient, WebClient> webClientConfigurer = Function.identity();
+	private Supplier<HttpHeaders> headersSupplier = () -> HttpHeaders.EMPTY;
 
 	/*
 	 * (non-Javadoc)
@@ -196,14 +198,24 @@ public TerminalClientConfigurationBuilder withPathPrefix(String pathPrefix) {
 	}
 
 	@Override
-	public TerminalClientConfigurationBuilder withWebClientConfigurer(Function<WebClient, WebClient> webClientConfigurer) {
+	public TerminalClientConfigurationBuilder withWebClientConfigurer(
+			Function<WebClient, WebClient> webClientConfigurer) {
 
 		Assert.notNull(webClientConfigurer, "webClientConfigurer must not be null");
 
 		this.webClientConfigurer = webClientConfigurer;
 		return this;
 	}
 
+	@Override
+	public TerminalClientConfigurationBuilder withHeaders(Supplier<HttpHeaders> headers) {
+
+		Assert.notNull(headers, "headersSupplier must not be null");
+
+		this.headersSupplier = headers;
+		return this;
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.elasticsearch.client.ClientConfiguration.ClientConfigurationBuilderWithOptionalDefaultHeaders#build()
@@ -219,7 +231,7 @@ public ClientConfiguration build() {
 		}
 
 		return new DefaultClientConfiguration(hosts, headers, useSsl, sslContext, soTimeout, connectTimeout, pathPrefix,
-				hostnameVerifier, proxy, webClientConfigurer);
+				hostnameVerifier, proxy, webClientConfigurer, headersSupplier);
 	}
 
 	private static InetSocketAddress parse(String hostAndPort) {

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

@@ -20,6 +20,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.http.HttpStatus;
+import org.springframework.lang.Nullable;
 import org.springframework.util.ObjectUtils;
 
 /**
@@ -90,7 +91,7 @@ public static void logRequest(String logId, String method, String endpoint, Obje
 	 * @param logId the correlation Id, see {@link #newLogId()}.
 	 * @param statusCode the HTTP status code.
 	 */
-	public static void logRawResponse(String logId, HttpStatus statusCode) {
+	public static void logRawResponse(String logId, @Nullable HttpStatus statusCode) {
 
 		if (isEnabled()) {
 			WIRE_LOGGER.trace("[{}] Received raw response: {}", logId, statusCode);

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

@@ -22,6 +22,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import javax.net.ssl.HostnameVerifier;
 import javax.net.ssl.SSLContext;
@@ -50,12 +51,13 @@ class DefaultClientConfiguration implements ClientConfiguration {
 	private final @Nullable String pathPrefix;
 	private final @Nullable HostnameVerifier hostnameVerifier;
 	private final @Nullable String proxy;
-	private final @Nullable Function<WebClient, WebClient> webClientConfigurer;
+	private final Function<WebClient, WebClient> webClientConfigurer;
+	private final Supplier<HttpHeaders> headersSupplier;
 
 	DefaultClientConfiguration(List<InetSocketAddress> hosts, HttpHeaders headers, boolean useSsl,
 			@Nullable SSLContext sslContext, Duration soTimeout, Duration connectTimeout, @Nullable String pathPrefix,
 			@Nullable HostnameVerifier hostnameVerifier, @Nullable String proxy,
-			@Nullable Function<WebClient, WebClient> webClientConfigurer) {
+			Function<WebClient, WebClient> webClientConfigurer, Supplier<HttpHeaders> headersSupplier) {
 
 		this.hosts = Collections.unmodifiableList(new ArrayList<>(hosts));
 		this.headers = new HttpHeaders(headers);
@@ -67,6 +69,7 @@ class DefaultClientConfiguration implements ClientConfiguration {
 		this.hostnameVerifier = hostnameVerifier;
 		this.proxy = proxy;
 		this.webClientConfigurer = webClientConfigurer;
+		this.headersSupplier = headersSupplier;
 	}
 
 	@Override
@@ -117,6 +120,11 @@ public Optional<String> getProxy() {
 
 	@Override
 	public Function<WebClient, WebClient> getWebClientConfigurer() {
-		return webClientConfigurer != null ? webClientConfigurer : Function.identity();
+		return webClientConfigurer;
+	}
+
+	@Override
+	public Supplier<HttpHeaders> getHeadersSupplier() {
+		return headersSupplier;
 	}
 }

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

@@ -20,13 +20,11 @@
 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.List;
-import java.util.Optional;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
-import javax.net.ssl.HostnameVerifier;
-import javax.net.ssl.SSLContext;
-
 import org.apache.http.Header;
 import org.apache.http.HttpEntity;
 import org.apache.http.HttpEntityEnclosingRequest;
@@ -87,38 +85,32 @@ public static ElasticsearchRestClient create(ClientConfiguration clientConfigura
 		HttpHeaders headers = clientConfiguration.getDefaultHeaders();
 
 		if (!headers.isEmpty()) {
-
-			Header[] httpHeaders = headers.toSingleValueMap().entrySet().stream()
-					.map(it -> new BasicHeader(it.getKey(), it.getValue())).toArray(Header[]::new);
-			builder.setDefaultHeaders(httpHeaders);
+			builder.setDefaultHeaders(toHeaderArray(headers));
 		}
 
 		builder.setHttpClientConfigCallback(clientBuilder -> {
-
-			Optional<SSLContext> sslContext = clientConfiguration.getSslContext();
-			Optional<HostnameVerifier> hostNameVerifier = clientConfiguration.getHostNameVerifier();
-			sslContext.ifPresent(clientBuilder::setSSLContext);
-			hostNameVerifier.ifPresent(clientBuilder::setSSLHostnameVerifier);
+			clientConfiguration.getSslContext().ifPresent(clientBuilder::setSSLContext);
+			clientConfiguration.getHostNameVerifier().ifPresent(clientBuilder::setSSLHostnameVerifier);
+			clientBuilder.addInterceptorLast(new CustomHeaderInjector(clientConfiguration.getHeadersSupplier()));
 
 			if (ClientLogger.isEnabled()) {
-
 				HttpLoggingInterceptor interceptor = new HttpLoggingInterceptor();
 
 				clientBuilder.addInterceptorLast((HttpRequestInterceptor) interceptor);
 				clientBuilder.addInterceptorLast((HttpResponseInterceptor) interceptor);
 			}
 
-			Duration connectTimeout = clientConfiguration.getConnectTimeout();
-			Duration timeout = clientConfiguration.getSocketTimeout();
-
 			Builder requestConfigBuilder = RequestConfig.custom();
+			Duration connectTimeout = clientConfiguration.getConnectTimeout();
 
 			if (!connectTimeout.isNegative()) {
 
 				requestConfigBuilder.setConnectTimeout(Math.toIntExact(connectTimeout.toMillis()));
 				requestConfigBuilder.setConnectionRequestTimeout(Math.toIntExact(connectTimeout.toMillis()));
 			}
 
+			Duration timeout = clientConfiguration.getSocketTimeout();
+
 			if (!timeout.isNegative()) {
 				requestConfigBuilder.setSocketTimeout(Math.toIntExact(timeout.toMillis()));
 			}
@@ -134,8 +126,16 @@ public static ElasticsearchRestClient create(ClientConfiguration clientConfigura
 		return () -> client;
 	}
 
+	private static Header[] toHeaderArray(HttpHeaders headers) {
+		return headers.entrySet().stream() //
+				.flatMap(entry -> entry.getValue().stream() //
+						.map(value -> new BasicHeader(entry.getKey(), value))) //
+				.toArray(Header[]::new);
+	}
+
 	private static List<String> formattedHosts(List<InetSocketAddress> hosts, boolean useSsl) {
-		return hosts.stream().map(it -> (useSsl ? "https" : "http") + "://" + it.getHostString() + ":" + it.getPort()).collect(Collectors.toList());
+		return hosts.stream().map(it -> (useSsl ? "https" : "http") + "://" + it.getHostString() + ":" + it.getPort())
+				.collect(Collectors.toList());
 	}
 
 	/**
@@ -180,7 +180,6 @@ public void process(HttpRequest request, HttpContext context) throws IOException
 			String logId = (String) context.getAttribute(RestClients.LOG_ID_ATTRIBUTE);
 
 			if (logId == null) {
-
 				logId = ClientLogger.newLogId();
 				context.setAttribute(RestClients.LOG_ID_ATTRIBUTE, logId);
 			}
@@ -205,10 +204,31 @@ public void process(HttpRequest request, HttpContext context) throws IOException
 
 		@Override
 		public void process(HttpResponse response, HttpContext context) {
-
 			String logId = (String) context.getAttribute(RestClients.LOG_ID_ATTRIBUTE);
-
 			ClientLogger.logRawResponse(logId, HttpStatus.resolve(response.getStatusLine().getStatusCode()));
 		}
 	}
+
+	/**
+	 * Interceptor to inject custom supplied headers.
+	 * 
+	 * @since 4.0
+	 */
+	private static class CustomHeaderInjector implements HttpRequestInterceptor {
+
+		public CustomHeaderInjector(Supplier<HttpHeaders> headersSupplier) {
+			this.headersSupplier = headersSupplier;
+		}
+
+		private final Supplier<HttpHeaders> headersSupplier;
+
+		@Override
+		public void process(HttpRequest request, HttpContext context) {
+			HttpHeaders httpHeaders = headersSupplier.get();
+
+			if (httpHeaders != null && httpHeaders != HttpHeaders.EMPTY) {
+				Arrays.stream(toHeaderArray(httpHeaders)).forEach(request::addHeader);
+			}
+		}
+	}
 }