Skip to content

Commit c21190d

Browse files
linarkounamsoo2
linarkou
authored and
namsoo2
committed
spring-projectsGH-2954: Opensearch vector store transport improvements
Fixes: spring-projects#2954 * fixed authorization to multiple hosts * ability to provide SSL certs via `SslBundles` * ability to set connectTimeout/readTimeout Signed-off-by: Linar Abzaltdinov <[email protected]> Signed-off-by: minsoo.nam <[email protected]>
1 parent 6ce5365 commit c21190d

File tree

4 files changed

+132
-20
lines changed

4 files changed

+132
-20
lines changed

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-opensearch/src/main/java/org/springframework/ai/vectorstore/opensearch/autoconfigure/OpenSearchVectorStoreAutoConfiguration.java‎

Lines changed: 65 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,19 @@
1717
package org.springframework.ai.vectorstore.opensearch.autoconfigure;
1818

1919
import java.net.URISyntaxException;
20+
import java.time.Duration;
2021
import java.util.List;
2122
import java.util.Optional;
23+
import java.util.concurrent.TimeUnit;
2224

2325
import io.micrometer.observation.ObservationRegistry;
2426
import org.apache.hc.client5.http.auth.AuthScope;
2527
import org.apache.hc.client5.http.auth.UsernamePasswordCredentials;
28+
import org.apache.hc.client5.http.config.RequestConfig;
2629
import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider;
30+
import org.apache.hc.client5.http.impl.nio.PoolingAsyncClientConnectionManagerBuilder;
31+
import org.apache.hc.client5.http.nio.AsyncClientConnectionManager;
32+
import org.apache.hc.client5.http.ssl.ClientTlsStrategyBuilder;
2733
import org.apache.hc.core5.http.HttpHost;
2834
import org.opensearch.client.opensearch.OpenSearchClient;
2935
import org.opensearch.client.transport.OpenSearchTransport;
@@ -33,7 +39,6 @@
3339
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
3440
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
3541
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
36-
import software.amazon.awssdk.http.SdkHttpClient;
3742
import software.amazon.awssdk.http.apache.ApacheHttpClient;
3843
import software.amazon.awssdk.regions.Region;
3944

@@ -50,6 +55,7 @@
5055
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingClass;
5156
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
5257
import org.springframework.boot.context.properties.EnableConfigurationProperties;
58+
import org.springframework.boot.ssl.SslBundles;
5359
import org.springframework.context.annotation.Bean;
5460
import org.springframework.context.annotation.Configuration;
5561
import org.springframework.util.StringUtils;
@@ -99,26 +105,57 @@ static class OpenSearchConfiguration {
99105

100106
@Bean
101107
@ConditionalOnMissingBean
102-
OpenSearchClient openSearchClient(OpenSearchConnectionDetails connectionDetails) {
103-
HttpHost[] httpHosts = connectionDetails.getUris()
104-
.stream()
105-
.map(s -> createHttpHost(s))
106-
.toArray(HttpHost[]::new);
107-
ApacheHttpClient5TransportBuilder transportBuilder = ApacheHttpClient5TransportBuilder.builder(httpHosts);
108-
Optional.ofNullable(connectionDetails.getUsername())
109-
.map(username -> createBasicCredentialsProvider(httpHosts[0], username,
110-
connectionDetails.getPassword()))
111-
.ifPresent(basicCredentialsProvider -> transportBuilder
112-
.setHttpClientConfigCallback(httpAsyncClientBuilder -> httpAsyncClientBuilder
113-
.setDefaultCredentialsProvider(basicCredentialsProvider)));
108+
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties, Optional<SslBundles> sslBundles) {
109+
HttpHost[] httpHosts = properties.getUris().stream().map(this::createHttpHost).toArray(HttpHost[]::new);
110+
Optional<BasicCredentialsProvider> basicCredentialsProvider = Optional.ofNullable(properties.getUsername())
111+
.map(username -> createBasicCredentialsProvider(httpHosts, username, properties.getPassword()));
112+
113+
var transportBuilder = ApacheHttpClient5TransportBuilder.builder(httpHosts);
114+
transportBuilder.setHttpClientConfigCallback(httpClientBuilder -> {
115+
basicCredentialsProvider.ifPresent(httpClientBuilder::setDefaultCredentialsProvider);
116+
httpClientBuilder.setConnectionManager(createConnectionManager(properties, sslBundles));
117+
httpClientBuilder.setDefaultRequestConfig(createRequestConfig(properties));
118+
return httpClientBuilder;
119+
});
120+
114121
return new OpenSearchClient(transportBuilder.build());
115122
}
116123

117-
private BasicCredentialsProvider createBasicCredentialsProvider(HttpHost httpHost, String username,
124+
private AsyncClientConnectionManager createConnectionManager(OpenSearchVectorStoreProperties properties,
125+
Optional<SslBundles> sslBundles) {
126+
var connectionManagerBuilder = PoolingAsyncClientConnectionManagerBuilder.create();
127+
if (sslBundles.isPresent()) {
128+
Optional.ofNullable(properties.getSslBundle())
129+
.map(bundle -> sslBundles.get().getBundle(bundle))
130+
.map(bundle -> ClientTlsStrategyBuilder.create()
131+
.setSslContext(bundle.createSslContext())
132+
.setTlsVersions(bundle.getOptions().getEnabledProtocols())
133+
.build())
134+
.ifPresent(connectionManagerBuilder::setTlsStrategy);
135+
}
136+
return connectionManagerBuilder.build();
137+
}
138+
139+
private RequestConfig createRequestConfig(OpenSearchVectorStoreProperties properties) {
140+
var requestConfigBuilder = RequestConfig.custom();
141+
Optional.ofNullable(properties.getConnectionTimeout())
142+
.map(Duration::toMillis)
143+
.ifPresent(timeoutMillis -> requestConfigBuilder.setConnectionRequestTimeout(timeoutMillis,
144+
TimeUnit.MILLISECONDS));
145+
Optional.ofNullable(properties.getReadTimeout())
146+
.map(Duration::toMillis)
147+
.ifPresent(
148+
timeoutMillis -> requestConfigBuilder.setResponseTimeout(timeoutMillis, TimeUnit.MILLISECONDS));
149+
return requestConfigBuilder.build();
150+
}
151+
152+
private BasicCredentialsProvider createBasicCredentialsProvider(HttpHost[] httpHosts, String username,
118153
String password) {
119154
BasicCredentialsProvider basicCredentialsProvider = new BasicCredentialsProvider();
120-
basicCredentialsProvider.setCredentials(new AuthScope(httpHost),
121-
new UsernamePasswordCredentials(username, password.toCharArray()));
155+
for (HttpHost httpHost : httpHosts) {
156+
basicCredentialsProvider.setCredentials(new AuthScope(httpHost),
157+
new UsernamePasswordCredentials(username, password.toCharArray()));
158+
}
122159
return basicCredentialsProvider;
123160
}
124161

@@ -159,12 +196,21 @@ PropertiesAwsOpenSearchConnectionDetails awsOpenSearchConnectionDetails(
159196

160197
@Bean
161198
@ConditionalOnMissingBean
162-
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties,
199+
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties, Optional<SslBundles> sslBundles,
163200
AwsOpenSearchConnectionDetails connectionDetails, AwsSdk2TransportOptions options) {
164201
Region region = Region.of(connectionDetails.getRegion());
165202

166-
SdkHttpClient httpClient = ApacheHttpClient.builder().build();
167-
OpenSearchTransport transport = new AwsSdk2Transport(httpClient,
203+
var httpClientBuilder = ApacheHttpClient.builder();
204+
Optional.ofNullable(properties.getConnectionTimeout()).ifPresent(httpClientBuilder::connectionTimeout);
205+
Optional.ofNullable(properties.getReadTimeout()).ifPresent(httpClientBuilder::socketTimeout);
206+
if (sslBundles.isPresent()) {
207+
Optional.ofNullable(properties.getSslBundle())
208+
.map(bundle -> sslBundles.get().getBundle(bundle))
209+
.ifPresent(bundle -> httpClientBuilder
210+
.tlsKeyManagersProvider(() -> bundle.getManagers().getKeyManagers())
211+
.tlsTrustManagersProvider(() -> bundle.getManagers().getTrustManagers()));
212+
}
213+
OpenSearchTransport transport = new AwsSdk2Transport(httpClientBuilder.build(),
168214
connectionDetails.getHost(properties.getAws().getDomainName()),
169215
properties.getAws().getServiceName(), region, options);
170216
return new OpenSearchClient(transport);

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-opensearch/src/main/java/org/springframework/ai/vectorstore/opensearch/autoconfigure/OpenSearchVectorStoreProperties.java‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 the original author or authors.
2+
* Copyright 2023-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
1616

1717
package org.springframework.ai.vectorstore.opensearch.autoconfigure;
1818

19+
import java.time.Duration;
1920
import java.util.List;
2021

2122
import org.springframework.ai.vectorstore.properties.CommonVectorStoreProperties;
@@ -39,6 +40,18 @@ public class OpenSearchVectorStoreProperties extends CommonVectorStoreProperties
3940

4041
private String mappingJson;
4142

43+
/**
44+
* SSL Bundle name ({@link org.springframework.boot.ssl.SslBundles}).
45+
*/
46+
private String sslBundle;
47+
48+
/**
49+
*
50+
*/
51+
private Duration connectionTimeout;
52+
53+
private Duration readTimeout;
54+
4255
private Aws aws = new Aws();
4356

4457
public List<String> getUris() {
@@ -81,6 +94,30 @@ public void setMappingJson(String mappingJson) {
8194
this.mappingJson = mappingJson;
8295
}
8396

97+
public String getSslBundle() {
98+
return sslBundle;
99+
}
100+
101+
public void setSslBundle(String sslBundle) {
102+
this.sslBundle = sslBundle;
103+
}
104+
105+
public Duration getConnectionTimeout() {
106+
return connectionTimeout;
107+
}
108+
109+
public void setConnectionTimeout(Duration connectionTimeout) {
110+
this.connectionTimeout = connectionTimeout;
111+
}
112+
113+
public Duration getReadTimeout() {
114+
return readTimeout;
115+
}
116+
117+
public void setReadTimeout(Duration readTimeout) {
118+
this.readTimeout = readTimeout;
119+
}
120+
84121
public Aws getAws() {
85122
return this.aws;
86123
}

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-opensearch/src/test/java/org/springframework/ai/vectorstore/opensearch/autoconfigure/AwsOpenSearchVectorStoreAutoConfigurationIT.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import org.awaitility.Awaitility;
2828
import org.junit.jupiter.api.BeforeAll;
2929
import org.junit.jupiter.api.Test;
30+
import org.opensearch.client.opensearch.OpenSearchClient;
3031
import org.testcontainers.containers.localstack.LocalStackContainer;
3132
import org.testcontainers.junit.jupiter.Container;
3233
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -37,8 +38,11 @@
3738
import org.springframework.ai.retry.autoconfigure.SpringAiRetryAutoConfiguration;
3839
import org.springframework.ai.transformers.TransformersEmbeddingModel;
3940
import org.springframework.ai.vectorstore.SearchRequest;
41+
import org.springframework.ai.vectorstore.VectorStore;
4042
import org.springframework.ai.vectorstore.opensearch.OpenSearchVectorStore;
4143
import org.springframework.boot.autoconfigure.AutoConfigurations;
44+
import org.springframework.boot.autoconfigure.ssl.SslAutoConfiguration;
45+
import org.springframework.boot.ssl.SslBundles;
4246
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
4347
import org.springframework.context.annotation.Bean;
4448
import org.springframework.context.annotation.Configuration;
@@ -138,6 +142,17 @@ public void addAndSearchTest() {
138142
});
139143
}
140144

145+
@Test
146+
public void autoConfigurationWithSslBundles() {
147+
this.contextRunner.withConfiguration(AutoConfigurations.of(SslAutoConfiguration.class)).run(context -> {
148+
assertThat(context.getBeansOfType(SslBundles.class)).isNotEmpty();
149+
assertThat(context.getBeansOfType(OpenSearchClient.class)).isNotEmpty();
150+
assertThat(context.getBeansOfType(OpenSearchVectorStoreProperties.class)).isNotEmpty();
151+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
152+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(OpenSearchVectorStore.class);
153+
});
154+
}
155+
141156
private String getText(String uri) {
142157
var resource = new DefaultResourceLoader().getResource(uri);
143158
try {

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-opensearch/src/test/java/org/springframework/ai/vectorstore/opensearch/autoconfigure/OpenSearchVectorStoreAutoConfigurationIT.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import io.micrometer.observation.tck.TestObservationRegistry;
2525
import org.awaitility.Awaitility;
2626
import org.junit.jupiter.api.Test;
27+
import org.opensearch.client.opensearch.OpenSearchClient;
2728
import org.opensearch.testcontainers.OpensearchContainer;
2829
import org.testcontainers.junit.jupiter.Container;
2930
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -42,6 +43,8 @@
4243
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
4344
import org.springframework.ai.vectorstore.opensearch.OpenSearchVectorStore;
4445
import org.springframework.boot.autoconfigure.AutoConfigurations;
46+
import org.springframework.boot.autoconfigure.ssl.SslAutoConfiguration;
47+
import org.springframework.boot.ssl.SslBundles;
4548
import org.springframework.boot.test.context.FilteredClassLoader;
4649
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
4750
import org.springframework.context.annotation.Bean;
@@ -171,6 +174,17 @@ public void autoConfigurationEnabledWhenTypeIsOpensearch() {
171174
});
172175
}
173176

177+
@Test
178+
public void autoConfigurationWithSslBundles() {
179+
this.contextRunner.withConfiguration(AutoConfigurations.of(SslAutoConfiguration.class)).run(context -> {
180+
assertThat(context.getBeansOfType(SslBundles.class)).isNotEmpty();
181+
assertThat(context.getBeansOfType(OpenSearchClient.class)).isNotEmpty();
182+
assertThat(context.getBeansOfType(OpenSearchVectorStoreProperties.class)).isNotEmpty();
183+
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
184+
assertThat(context.getBean(VectorStore.class)).isInstanceOf(OpenSearchVectorStore.class);
185+
});
186+
}
187+
174188
private String getText(String uri) {
175189
var resource = new DefaultResourceLoader().getResource(uri);
176190
try {

0 commit comments

Comments
 (0)