Improve Elasticsearch RestClient customization capabilities

At present, RestClientBuilderCustomizer allows general customization of RestClientBuilder.
This is troublesome for users that want to customize `HttpAsyncClientBuilder` and
`RequestConfig.Builder` since those are set on the `RestClientBuilder`. By customizing
those two builders user lose out on Spring Boot's support for binding username, password,
connection-timeout and read-timeout properties from `"spring.elasticsearch.rest"` namespace.

This commit enhances the `RestClientBuilderCustomizer` with support for customizing
`HttpAsyncClientBuilder` and `RequestConfig.Builder` by providing additional `customize`
methods that accept the aforementioned builders. Both new methods are optional as they have
no-op default implementations.

See gh-20994
pull/21069/head
Vedran Pavic 5 years ago committed by Brian Clozel
parent 8de0027757
commit f701d97b92

@ -23,7 +23,9 @@ import org.apache.http.auth.AuthScope;
import org.apache.http.auth.Credentials;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
@ -40,32 +42,30 @@ import org.springframework.context.annotation.Configuration;
*
* @author Brian Clozel
* @author Stephane Nicoll
* @author Vedran Pavic
*/
class ElasticsearchRestClientConfigurations {
@Configuration(proxyBeanMethods = false)
@ConditionalOnMissingBean(RestClientBuilder.class)
static class RestClientBuilderConfiguration {
@Bean
@ConditionalOnMissingBean
RestClientBuilderCustomizer defaultRestClientBuilderCustomizer(ElasticsearchRestClientProperties properties) {
return new DefaultRestClientBuilderCustomizer(properties);
}
@Bean
RestClientBuilder elasticsearchRestClientBuilder(ElasticsearchRestClientProperties properties,
ObjectProvider<RestClientBuilderCustomizer> builderCustomizers) {
HttpHost[] hosts = properties.getUris().stream().map(HttpHost::create).toArray(HttpHost[]::new);
RestClientBuilder builder = RestClient.builder(hosts);
PropertyMapper map = PropertyMapper.get();
map.from(properties::getUsername).whenHasText().to((username) -> {
CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
Credentials credentials = new UsernamePasswordCredentials(properties.getUsername(),
properties.getPassword());
credentialsProvider.setCredentials(AuthScope.ANY, credentials);
builder.setHttpClientConfigCallback(
(httpClientBuilder) -> httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider));
builder.setHttpClientConfigCallback((httpClientBuilder) -> {
builderCustomizers.orderedStream().forEach((customizer) -> customizer.customize(httpClientBuilder));
return httpClientBuilder;
});
builder.setRequestConfigCallback((requestConfigBuilder) -> {
map.from(properties::getConnectionTimeout).whenNonNull().asInt(Duration::toMillis)
.to(requestConfigBuilder::setConnectTimeout);
map.from(properties::getReadTimeout).whenNonNull().asInt(Duration::toMillis)
.to(requestConfigBuilder::setSocketTimeout);
builderCustomizers.orderedStream().forEach((customizer) -> customizer.customize(requestConfigBuilder));
return requestConfigBuilder;
});
builderCustomizers.orderedStream().forEach((customizer) -> customizer.customize(builder));
@ -108,4 +108,39 @@ class ElasticsearchRestClientConfigurations {
}
static class DefaultRestClientBuilderCustomizer implements RestClientBuilderCustomizer {
private static final PropertyMapper map = PropertyMapper.get();
private final ElasticsearchRestClientProperties properties;
DefaultRestClientBuilderCustomizer(ElasticsearchRestClientProperties properties) {
this.properties = properties;
}
@Override
public void customize(RestClientBuilder builder) {
}
@Override
public void customize(HttpAsyncClientBuilder builder) {
map.from(this.properties::getUsername).whenHasText().to((username) -> {
CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
Credentials credentials = new UsernamePasswordCredentials(this.properties.getUsername(),
this.properties.getPassword());
credentialsProvider.setCredentials(AuthScope.ANY, credentials);
builder.setDefaultCredentialsProvider(credentialsProvider);
});
}
@Override
public void customize(RequestConfig.Builder builder) {
map.from(this.properties::getConnectionTimeout).whenNonNull().asInt(Duration::toMillis)
.to(builder::setConnectTimeout);
map.from(this.properties::getReadTimeout).whenNonNull().asInt(Duration::toMillis)
.to(builder::setSocketTimeout);
}
}
}

@ -16,6 +16,8 @@
package org.springframework.boot.autoconfigure.elasticsearch;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
import org.elasticsearch.client.RestClientBuilder;
/**
@ -24,6 +26,7 @@ import org.elasticsearch.client.RestClientBuilder;
* retaining default auto-configuration.
*
* @author Brian Clozel
* @author Vedran Pavic
* @since 2.1.0
*/
@FunctionalInterface
@ -35,4 +38,18 @@ public interface RestClientBuilderCustomizer {
*/
void customize(RestClientBuilder builder);
/**
* Customize the {@link HttpAsyncClientBuilder}.
* @param builder the builder
*/
default void customize(HttpAsyncClientBuilder builder) {
}
/**
* Customize the {@link RequestConfig.Builder}.
* @param builder the builder
*/
default void customize(RequestConfig.Builder builder) {
}
}

@ -20,6 +20,8 @@ import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.RequestOptions;
@ -36,7 +38,6 @@ import org.springframework.boot.test.context.FilteredClassLoader;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.test.util.ReflectionTestUtils;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.mock;
@ -45,6 +46,7 @@ import static org.mockito.Mockito.mock;
* Tests for {@link ElasticsearchRestClientAutoConfiguration}.
*
* @author Brian Clozel
* @author Vedran Pavic
*/
@Testcontainers(disabledWithoutDocker = true)
class ElasticsearchRestClientAutoConfigurationTests {
@ -53,7 +55,7 @@ class ElasticsearchRestClientAutoConfigurationTests {
static final ElasticsearchContainer elasticsearch = new ElasticsearchContainer().withStartupAttempts(5)
.withStartupTimeout(Duration.ofMinutes(10));
private ApplicationContextRunner contextRunner = new ApplicationContextRunner()
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(ElasticsearchRestClientAutoConfiguration.class));
@Test
@ -106,6 +108,8 @@ class ElasticsearchRestClientAutoConfigurationTests {
assertThat(context).hasSingleBean(RestClient.class);
RestClient restClient = context.getBean(RestClient.class);
assertThat(restClient).hasFieldOrPropertyWithValue("pathPrefix", "/test");
assertThat(restClient).extracting("client.connmgr.pool.maxTotal").isEqualTo(100);
assertThat(restClient).extracting("client.defaultConfig.cookieSpec").isEqualTo("rfc6265-lax");
});
}
@ -130,10 +134,10 @@ class ElasticsearchRestClientAutoConfigurationTests {
}
private static void assertTimeouts(RestClient restClient, Duration connectTimeout, Duration readTimeout) {
Object client = ReflectionTestUtils.getField(restClient, "client");
Object config = ReflectionTestUtils.getField(client, "defaultConfig");
assertThat(config).hasFieldOrPropertyWithValue("socketTimeout", Math.toIntExact(readTimeout.toMillis()));
assertThat(config).hasFieldOrPropertyWithValue("connectTimeout", Math.toIntExact(connectTimeout.toMillis()));
assertThat(restClient).extracting("client.defaultConfig.socketTimeout")
.isEqualTo(Math.toIntExact(readTimeout.toMillis()));
assertThat(restClient).extracting("client.defaultConfig.connectTimeout")
.isEqualTo(Math.toIntExact(connectTimeout.toMillis()));
}
@Test
@ -167,7 +171,24 @@ class ElasticsearchRestClientAutoConfigurationTests {
@Bean
RestClientBuilderCustomizer myCustomizer() {
return (builder) -> builder.setPathPrefix("/test");
return new RestClientBuilderCustomizer() {
@Override
public void customize(RestClientBuilder builder) {
builder.setPathPrefix("/test");
}
@Override
public void customize(HttpAsyncClientBuilder builder) {
builder.setMaxConnTotal(100);
}
@Override
public void customize(RequestConfig.Builder builder) {
builder.setCookieSpec("rfc6265-lax");
}
};
}
}

Loading…
Cancel
Save