Skip to content

Commit f4eaa4d

Browse files
linarkoulinarkou
committed
vector-store-opensearch:
* fixed authorization to multiple hosts * ability to provide SSL certs via `SslBundles` * ability to set connectTimeout/readTimeout Signed-off-by: Linar Abzaltdinov <[email protected]>
1 parent 88490b3 commit f4eaa4d

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;
@@ -100,26 +106,57 @@ static class OpenSearchConfiguration {
100106

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

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

@@ -147,12 +184,21 @@ PropertiesAwsOpenSearchConnectionDetails awsOpenSearchConnectionDetails(
147184

148185
@Bean
149186
@ConditionalOnMissingBean
150-
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties,
187+
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties, Optional<SslBundles> sslBundles,
151188
AwsOpenSearchConnectionDetails connectionDetails, AwsSdk2TransportOptions options) {
152189
Region region = Region.of(connectionDetails.getRegion());
153190

154-
SdkHttpClient httpClient = ApacheHttpClient.builder().build();
155-
OpenSearchTransport transport = new AwsSdk2Transport(httpClient,
191+
var httpClientBuilder = ApacheHttpClient.builder();
192+
Optional.ofNullable(properties.getConnectionTimeout()).ifPresent(httpClientBuilder::connectionTimeout);
193+
Optional.ofNullable(properties.getReadTimeout()).ifPresent(httpClientBuilder::socketTimeout);
194+
if (sslBundles.isPresent()) {
195+
Optional.ofNullable(properties.getSslBundle())
196+
.map(bundle -> sslBundles.get().getBundle(bundle))
197+
.ifPresent(bundle -> httpClientBuilder
198+
.tlsKeyManagersProvider(() -> bundle.getManagers().getKeyManagers())
199+
.tlsTrustManagersProvider(() -> bundle.getManagers().getTrustManagers()));
200+
}
201+
OpenSearchTransport transport = new AwsSdk2Transport(httpClientBuilder.build(),
156202
connectionDetails.getHost(properties.getAws().getDomainName()),
157203
properties.getAws().getServiceName(), region, options);
158204
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;
@@ -170,6 +173,17 @@ public void autoConfigurationEnabledWhenTypeIsOpensearch() {
170173
});
171174
}
172175

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

0 commit comments

Comments
 (0)