Skip to content

Commit 11c1980

Browse files
committed
Merge pull request #20994 from vpavic
* improve-es-config: Polish Improve Elasticsearch RestClient customization capabilities Closes gh-20994
2 parents 8de0027 + f103966 commit 11c1980

File tree

3 files changed

+100
-20
lines changed

3 files changed

+100
-20
lines changed

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/elasticsearch/ElasticsearchRestClientConfigurations.java

Lines changed: 48 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@
2323
import org.apache.http.auth.Credentials;
2424
import org.apache.http.auth.UsernamePasswordCredentials;
2525
import org.apache.http.client.CredentialsProvider;
26+
import org.apache.http.client.config.RequestConfig;
2627
import org.apache.http.impl.client.BasicCredentialsProvider;
28+
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
2729
import org.elasticsearch.client.RestClient;
2830
import org.elasticsearch.client.RestClientBuilder;
2931
import org.elasticsearch.client.RestHighLevelClient;
@@ -40,32 +42,30 @@
4042
*
4143
* @author Brian Clozel
4244
* @author Stephane Nicoll
45+
* @author Vedran Pavic
4346
*/
4447
class ElasticsearchRestClientConfigurations {
4548

4649
@Configuration(proxyBeanMethods = false)
50+
@ConditionalOnMissingBean(RestClientBuilder.class)
4751
static class RestClientBuilderConfiguration {
4852

4953
@Bean
50-
@ConditionalOnMissingBean
54+
RestClientBuilderCustomizer defaultRestClientBuilderCustomizer(ElasticsearchRestClientProperties properties) {
55+
return new DefaultRestClientBuilderCustomizer(properties);
56+
}
57+
58+
@Bean
5159
RestClientBuilder elasticsearchRestClientBuilder(ElasticsearchRestClientProperties properties,
5260
ObjectProvider<RestClientBuilderCustomizer> builderCustomizers) {
5361
HttpHost[] hosts = properties.getUris().stream().map(HttpHost::create).toArray(HttpHost[]::new);
5462
RestClientBuilder builder = RestClient.builder(hosts);
55-
PropertyMapper map = PropertyMapper.get();
56-
map.from(properties::getUsername).whenHasText().to((username) -> {
57-
CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
58-
Credentials credentials = new UsernamePasswordCredentials(properties.getUsername(),
59-
properties.getPassword());
60-
credentialsProvider.setCredentials(AuthScope.ANY, credentials);
61-
builder.setHttpClientConfigCallback(
62-
(httpClientBuilder) -> httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider));
63+
builder.setHttpClientConfigCallback((httpClientBuilder) -> {
64+
builderCustomizers.orderedStream().forEach((customizer) -> customizer.customize(httpClientBuilder));
65+
return httpClientBuilder;
6366
});
6467
builder.setRequestConfigCallback((requestConfigBuilder) -> {
65-
map.from(properties::getConnectionTimeout).whenNonNull().asInt(Duration::toMillis)
66-
.to(requestConfigBuilder::setConnectTimeout);
67-
map.from(properties::getReadTimeout).whenNonNull().asInt(Duration::toMillis)
68-
.to(requestConfigBuilder::setSocketTimeout);
68+
builderCustomizers.orderedStream().forEach((customizer) -> customizer.customize(requestConfigBuilder));
6969
return requestConfigBuilder;
7070
});
7171
builderCustomizers.orderedStream().forEach((customizer) -> customizer.customize(builder));
@@ -108,4 +108,39 @@ RestClient elasticsearchRestClient(RestClientBuilder builder) {
108108

109109
}
110110

111+
static class DefaultRestClientBuilderCustomizer implements RestClientBuilderCustomizer {
112+
113+
private static final PropertyMapper map = PropertyMapper.get();
114+
115+
private final ElasticsearchRestClientProperties properties;
116+
117+
DefaultRestClientBuilderCustomizer(ElasticsearchRestClientProperties properties) {
118+
this.properties = properties;
119+
}
120+
121+
@Override
122+
public void customize(RestClientBuilder builder) {
123+
}
124+
125+
@Override
126+
public void customize(HttpAsyncClientBuilder builder) {
127+
map.from(this.properties::getUsername).whenHasText().to((username) -> {
128+
CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
129+
Credentials credentials = new UsernamePasswordCredentials(this.properties.getUsername(),
130+
this.properties.getPassword());
131+
credentialsProvider.setCredentials(AuthScope.ANY, credentials);
132+
builder.setDefaultCredentialsProvider(credentialsProvider);
133+
});
134+
}
135+
136+
@Override
137+
public void customize(RequestConfig.Builder builder) {
138+
map.from(this.properties::getConnectionTimeout).whenNonNull().asInt(Duration::toMillis)
139+
.to(builder::setConnectTimeout);
140+
map.from(this.properties::getReadTimeout).whenNonNull().asInt(Duration::toMillis)
141+
.to(builder::setSocketTimeout);
142+
}
143+
144+
}
145+
111146
}

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/elasticsearch/RestClientBuilderCustomizer.java

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
package org.springframework.boot.autoconfigure.elasticsearch;
1818

19+
import org.apache.http.client.config.RequestConfig;
20+
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
1921
import org.elasticsearch.client.RestClientBuilder;
2022

2123
/**
@@ -24,15 +26,37 @@
2426
* retaining default auto-configuration.
2527
*
2628
* @author Brian Clozel
29+
* @author Vedran Pavic
2730
* @since 2.1.0
2831
*/
2932
@FunctionalInterface
3033
public interface RestClientBuilderCustomizer {
3134

3235
/**
3336
* Customize the {@link RestClientBuilder}.
37+
* <p>
38+
* Possibly overrides customizations made with the {@code "spring.elasticsearch.rest"}
39+
* configuration properties namespace. For more targeted changes, see
40+
* {@link #customize(HttpAsyncClientBuilder)} and
41+
* {@link #customize(RequestConfig.Builder)}.
3442
* @param builder the builder to customize
3543
*/
3644
void customize(RestClientBuilder builder);
3745

46+
/**
47+
* Customize the {@link HttpAsyncClientBuilder}.
48+
* @param builder the builder
49+
* @since 2.3.0
50+
*/
51+
default void customize(HttpAsyncClientBuilder builder) {
52+
}
53+
54+
/**
55+
* Customize the {@link RequestConfig.Builder}.
56+
* @param builder the builder
57+
* @since 2.3.0
58+
*/
59+
default void customize(RequestConfig.Builder builder) {
60+
}
61+
3862
}

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/elasticsearch/ElasticsearchRestClientAutoConfigurationTests.java

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
import java.util.HashMap;
2121
import java.util.Map;
2222

23+
import org.apache.http.client.config.RequestConfig;
24+
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
2325
import org.elasticsearch.action.get.GetRequest;
2426
import org.elasticsearch.action.index.IndexRequest;
2527
import org.elasticsearch.client.RequestOptions;
@@ -36,7 +38,6 @@
3638
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
3739
import org.springframework.context.annotation.Bean;
3840
import org.springframework.context.annotation.Configuration;
39-
import org.springframework.test.util.ReflectionTestUtils;
4041

4142
import static org.assertj.core.api.Assertions.assertThat;
4243
import static org.mockito.Mockito.mock;
@@ -45,6 +46,7 @@
4546
* Tests for {@link ElasticsearchRestClientAutoConfiguration}.
4647
*
4748
* @author Brian Clozel
49+
* @author Vedran Pavic
4850
*/
4951
@Testcontainers(disabledWithoutDocker = true)
5052
class ElasticsearchRestClientAutoConfigurationTests {
@@ -53,7 +55,7 @@ class ElasticsearchRestClientAutoConfigurationTests {
5355
static final ElasticsearchContainer elasticsearch = new ElasticsearchContainer().withStartupAttempts(5)
5456
.withStartupTimeout(Duration.ofMinutes(10));
5557

56-
private ApplicationContextRunner contextRunner = new ApplicationContextRunner()
58+
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
5759
.withConfiguration(AutoConfigurations.of(ElasticsearchRestClientAutoConfiguration.class));
5860

5961
@Test
@@ -106,6 +108,8 @@ void configureWhenBuilderCustomizerShouldApply() {
106108
assertThat(context).hasSingleBean(RestClient.class);
107109
RestClient restClient = context.getBean(RestClient.class);
108110
assertThat(restClient).hasFieldOrPropertyWithValue("pathPrefix", "/test");
111+
assertThat(restClient).extracting("client.connmgr.pool.maxTotal").isEqualTo(100);
112+
assertThat(restClient).extracting("client.defaultConfig.cookieSpec").isEqualTo("rfc6265-lax");
109113
});
110114
}
111115

@@ -130,10 +134,10 @@ void configureWithCustomTimeouts() {
130134
}
131135

132136
private static void assertTimeouts(RestClient restClient, Duration connectTimeout, Duration readTimeout) {
133-
Object client = ReflectionTestUtils.getField(restClient, "client");
134-
Object config = ReflectionTestUtils.getField(client, "defaultConfig");
135-
assertThat(config).hasFieldOrPropertyWithValue("socketTimeout", Math.toIntExact(readTimeout.toMillis()));
136-
assertThat(config).hasFieldOrPropertyWithValue("connectTimeout", Math.toIntExact(connectTimeout.toMillis()));
137+
assertThat(restClient).extracting("client.defaultConfig.socketTimeout")
138+
.isEqualTo(Math.toIntExact(readTimeout.toMillis()));
139+
assertThat(restClient).extracting("client.defaultConfig.connectTimeout")
140+
.isEqualTo(Math.toIntExact(connectTimeout.toMillis()));
137141
}
138142

139143
@Test
@@ -167,7 +171,24 @@ static class BuilderCustomizerConfiguration {
167171

168172
@Bean
169173
RestClientBuilderCustomizer myCustomizer() {
170-
return (builder) -> builder.setPathPrefix("/test");
174+
return new RestClientBuilderCustomizer() {
175+
176+
@Override
177+
public void customize(RestClientBuilder builder) {
178+
builder.setPathPrefix("/test");
179+
}
180+
181+
@Override
182+
public void customize(HttpAsyncClientBuilder builder) {
183+
builder.setMaxConnTotal(100);
184+
}
185+
186+
@Override
187+
public void customize(RequestConfig.Builder builder) {
188+
builder.setCookieSpec("rfc6265-lax");
189+
}
190+
191+
};
171192
}
172193

173194
}

0 commit comments

Comments
 (0)