Skip to content

Commit 6a5c7a5

Browse files
eddumelendezeddumelendez
committed
Add LocalStack OpenSearch Service Connection support for Docker Compose and Testcontainers
* Add property `spring.ai.vectorstore.opensearch.aws.domain-name` * Require `AwsCredentialsProvider` to enable `AwsOpenSearchConfiguration` * Add Testcontainers Service Connection support * Add Docker Compose Service Connection support
1 parent 420bf5e commit 6a5c7a5

File tree

16 files changed

+584
-63
lines changed

16 files changed

+584
-63
lines changed

‎spring-ai-docs/src/main/antora/modules/ROOT/pages/api/docker-compose.adoc‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@ The following service connection factories are provided in the `spring-ai-spring
3131
[cols="|,|"]
3232
|====
3333
| Connection Details | Matched on
34+
| `AwsOpenSearchConnectionDetails`
35+
| Containers named `localstack/localstack`
36+
3437
| `ChromaConnectionDetails`
3538
| Containers named `chromadb/chroma`, `ghcr.io/chroma-core/chroma`
3639

‎spring-ai-docs/src/main/antora/modules/ROOT/pages/api/testcontainers.adoc‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@ The following service connection factories are provided in the `spring-ai-spring
3131
[cols="|,|"]
3232
|====
3333
| Connection Details | Matched on
34+
35+
| `AwsOpenSearchConnectionDetails`
36+
| Containers of type `LocalStackContainer`
37+
3438
| `ChromaConnectionDetails`
3539
| Containers of type `ChromaDBContainer`
3640

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/opensearch/AwsOpenSearchConnectionDetails.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
/*
2+
* Copyright 2023-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.autoconfigure.vectorstore.opensearch;
18+
19+
import org.springframework.boot.autoconfigure.service.connection.ConnectionDetails;
20+
21+
public interface AwsOpenSearchConnectionDetails extends ConnectionDetails {
22+
23+
String getRegion();
24+
25+
String getAccessKey();
26+
27+
String getSecretKey();
28+
29+
String getHost(String domainName);
30+
31+
}

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/opensearch/OpenSearchVectorStoreAutoConfiguration.java‎

Lines changed: 52 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,9 @@
3030
import org.opensearch.client.transport.aws.AwsSdk2Transport;
3131
import org.opensearch.client.transport.aws.AwsSdk2TransportOptions;
3232
import org.opensearch.client.transport.httpclient5.ApacheHttpClient5TransportBuilder;
33+
import org.springframework.util.StringUtils;
3334
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
35+
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
3436
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
3537
import software.amazon.awssdk.http.SdkHttpClient;
3638
import software.amazon.awssdk.http.apache.ApacheHttpClient;
@@ -123,28 +125,35 @@ private HttpHost createHttpHost(String s) {
123125
}
124126

125127
@Configuration(proxyBeanMethods = false)
126-
@ConditionalOnClass({ Region.class, ApacheHttpClient.class })
128+
@ConditionalOnClass({ AwsCredentialsProvider.class, Region.class, ApacheHttpClient.class })
127129
static class AwsOpenSearchConfiguration {
128130

131+
@Bean
132+
@ConditionalOnMissingBean(AwsOpenSearchConnectionDetails.class)
133+
PropertiesAwsOpenSearchConnectionDetails awsOpenSearchConnectionDetails(
134+
OpenSearchVectorStoreProperties properties) {
135+
return new PropertiesAwsOpenSearchConnectionDetails(properties);
136+
}
137+
129138
@Bean
130139
@ConditionalOnMissingBean
131-
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties, AwsSdk2TransportOptions options) {
132-
OpenSearchVectorStoreProperties.Aws aws = properties.getAws();
133-
Region region = Region.of(aws.getRegion());
140+
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties,
141+
AwsOpenSearchConnectionDetails connectionDetails, AwsSdk2TransportOptions options) {
142+
Region region = Region.of(connectionDetails.getRegion());
134143

135144
SdkHttpClient httpClient = ApacheHttpClient.builder().build();
136-
OpenSearchTransport transport = new AwsSdk2Transport(httpClient, aws.getHost(), aws.getServiceName(),
137-
region, options);
145+
OpenSearchTransport transport = new AwsSdk2Transport(httpClient,
146+
connectionDetails.getHost(properties.getAws().getDomainName()),
147+
properties.getAws().getServiceName(), region, options);
138148
return new OpenSearchClient(transport);
139149
}
140150

141151
@Bean
142152
@ConditionalOnMissingBean
143-
AwsSdk2TransportOptions options(OpenSearchVectorStoreProperties properties) {
144-
OpenSearchVectorStoreProperties.Aws aws = properties.getAws();
153+
AwsSdk2TransportOptions options(AwsOpenSearchConnectionDetails connectionDetails) {
145154
return AwsSdk2TransportOptions.builder()
146-
.setCredentials(StaticCredentialsProvider
147-
.create(AwsBasicCredentials.create(aws.getAccessKey(), aws.getSecretKey())))
155+
.setCredentials(StaticCredentialsProvider.create(
156+
AwsBasicCredentials.create(connectionDetails.getAccessKey(), connectionDetails.getSecretKey())))
148157
.build();
149158
}
150159

@@ -175,4 +184,37 @@ public String getPassword() {
175184

176185
}
177186

187+
static class PropertiesAwsOpenSearchConnectionDetails implements AwsOpenSearchConnectionDetails {
188+
189+
private final OpenSearchVectorStoreProperties.Aws aws;
190+
191+
public PropertiesAwsOpenSearchConnectionDetails(OpenSearchVectorStoreProperties properties) {
192+
this.aws = properties.getAws();
193+
}
194+
195+
@Override
196+
public String getRegion() {
197+
return this.aws.getRegion();
198+
}
199+
200+
@Override
201+
public String getAccessKey() {
202+
return this.aws.getAccessKey();
203+
}
204+
205+
@Override
206+
public String getSecretKey() {
207+
return this.aws.getSecretKey();
208+
}
209+
210+
@Override
211+
public String getHost(String domainName) {
212+
if (StringUtils.hasText(domainName)) {
213+
return "%s.%s".formatted(this.aws.getDomainName(), this.aws.getHost());
214+
}
215+
return this.aws.getHost();
216+
}
217+
218+
}
219+
178220
}

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/opensearch/OpenSearchVectorStoreProperties.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,8 @@ public void setAws(Aws aws) {
9191

9292
static class Aws {
9393

94+
private String domainName;
95+
9496
private String host;
9597

9698
private String serviceName;
@@ -101,6 +103,14 @@ static class Aws {
101103

102104
private String region;
103105

106+
public String getDomainName() {
107+
return this.domainName;
108+
}
109+
110+
public void setDomainName(String domainName) {
111+
this.domainName = domainName;
112+
}
113+
104114
public String getHost() {
105115
return this.host;
106116
}

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/opensearch/AwsOpenSearchVectorStoreAutoConfigurationIT.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,12 +62,11 @@ class AwsOpenSearchVectorStoreAutoConfigurationIT {
6262
.withConfiguration(AutoConfigurations.of(OpenSearchVectorStoreAutoConfiguration.class,
6363
SpringAiRetryAutoConfiguration.class))
6464
.withUserConfiguration(Config.class)
65-
.withPropertyValues("spring.ai.vectorstore.opensearch.initialize-schema=true")
66-
.withPropertyValues(
65+
.withPropertyValues("spring.ai.vectorstore.opensearch.initialize-schema=true",
6766
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.host="
6867
+ String.format("testcontainers-domain.%s.opensearch.localhost.localstack.cloud:%s",
6968
localstack.getRegion(), localstack.getMappedPort(4566)),
70-
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.service-name=opensearch",
69+
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.service-name=es",
7170
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.region=" + localstack.getRegion(),
7271
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.access-key=" + localstack.getAccessKey(),
7372
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.secret-key=" + localstack.getSecretKey(),

‎spring-ai-spring-boot-docker-compose/src/main/java/org/springframework/ai/docker/compose/service/connection/opensearch/AwsOpenSearchDockerComposeConnectionDetailsFactory.java‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
/*
2+
* Copyright 2023-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.docker.compose.service.connection.opensearch;
18+
19+
import org.springframework.ai.autoconfigure.vectorstore.opensearch.AwsOpenSearchConnectionDetails;
20+
import org.springframework.ai.autoconfigure.vectorstore.opensearch.OpenSearchConnectionDetails;
21+
import org.springframework.boot.docker.compose.core.RunningService;
22+
import org.springframework.boot.docker.compose.service.connection.DockerComposeConnectionDetailsFactory;
23+
import org.springframework.boot.docker.compose.service.connection.DockerComposeConnectionSource;
24+
25+
/**
26+
* @author Eddú Meléndez
27+
*/
28+
class AwsOpenSearchDockerComposeConnectionDetailsFactory
29+
extends DockerComposeConnectionDetailsFactory<AwsOpenSearchConnectionDetails> {
30+
31+
private static final int LOCALSTACK_PORT = 4566;
32+
33+
protected AwsOpenSearchDockerComposeConnectionDetailsFactory() {
34+
super("localstack/localstack");
35+
}
36+
37+
@Override
38+
protected AwsOpenSearchConnectionDetails getDockerComposeConnectionDetails(DockerComposeConnectionSource source) {
39+
return new AwsOpenSearchDockerComposeConnectionDetails(source.getRunningService());
40+
}
41+
42+
/**
43+
* {@link OpenSearchConnectionDetails} backed by a {@code OpenSearch}
44+
* {@link RunningService}.
45+
*/
46+
static class AwsOpenSearchDockerComposeConnectionDetails extends DockerComposeConnectionDetails
47+
implements AwsOpenSearchConnectionDetails {
48+
49+
private final AwsOpenSearchEnvironment environment;
50+
51+
private final int port;
52+
53+
AwsOpenSearchDockerComposeConnectionDetails(RunningService service) {
54+
super(service);
55+
this.environment = new AwsOpenSearchEnvironment(service.env());
56+
this.port = service.ports().get(LOCALSTACK_PORT);
57+
}
58+
59+
@Override
60+
public String getRegion() {
61+
return this.environment.getRegion();
62+
}
63+
64+
@Override
65+
public String getAccessKey() {
66+
return this.environment.getAccessKey();
67+
}
68+
69+
@Override
70+
public String getSecretKey() {
71+
return this.environment.getSecretKey();
72+
}
73+
74+
@Override
75+
public String getHost(String domainName) {
76+
return "%s.%s.opensearch.localhost.localstack.cloud:%s".formatted(domainName, this.environment.getRegion(),
77+
this.port);
78+
}
79+
80+
}
81+
82+
}

‎spring-ai-spring-boot-docker-compose/src/main/java/org/springframework/ai/docker/compose/service/connection/opensearch/AwsOpenSearchEnvironment.java‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/*
2+
* Copyright 2023-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.docker.compose.service.connection.opensearch;
18+
19+
import java.util.Map;
20+
21+
class AwsOpenSearchEnvironment {
22+
23+
private final String region;
24+
25+
private final String accessKey;
26+
27+
private final String secretKey;
28+
29+
AwsOpenSearchEnvironment(Map<String, String> env) {
30+
this.region = env.getOrDefault("DEFAULT_REGION", "us-east-1");
31+
this.accessKey = env.getOrDefault("AWS_ACCESS_KEY_ID", "test");
32+
this.secretKey = env.getOrDefault("AWS_SECRET_ACCESS_KEY", "test");
33+
}
34+
35+
public String getRegion() {
36+
return this.region;
37+
}
38+
39+
public String getAccessKey() {
40+
return this.accessKey;
41+
}
42+
43+
public String getSecretKey() {
44+
return this.secretKey;
45+
}
46+
47+
}

‎spring-ai-spring-boot-docker-compose/src/main/resources/META-INF/spring.factories‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ org.springframework.boot.autoconfigure.service.connection.ConnectionDetailsFacto
1818
org.springframework.ai.docker.compose.service.connection.chroma.ChromaDockerComposeConnectionDetailsFactory,\
1919
org.springframework.ai.docker.compose.service.connection.mongo.MongoDbAtlasLocalDockerComposeConnectionDetailsFactory,\
2020
org.springframework.ai.docker.compose.service.connection.ollama.OllamaDockerComposeConnectionDetailsFactory,\
21+
org.springframework.ai.docker.compose.service.connection.opensearch.AwsOpenSearchDockerComposeConnectionDetailsFactory,\
2122
org.springframework.ai.docker.compose.service.connection.opensearch.OpenSearchDockerComposeConnectionDetailsFactory,\
2223
org.springframework.ai.docker.compose.service.connection.qdrant.QdrantDockerComposeConnectionDetailsFactory,\
2324
org.springframework.ai.docker.compose.service.connection.typesense.TypesenseDockerComposeConnectionDetailsFactory,\

‎spring-ai-spring-boot-docker-compose/src/test/java/org/springframework/ai/docker/compose/service/connection/opensearch/AwsOpenSearchDockerComposeConnectionDetailsFactoryTests.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
* Copyright 2023-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.docker.compose.service.connection.opensearch;
18+
19+
import org.junit.jupiter.api.Test;
20+
import org.springframework.ai.autoconfigure.vectorstore.opensearch.AwsOpenSearchConnectionDetails;
21+
import org.springframework.boot.docker.compose.service.connection.test.AbstractDockerComposeIntegrationTests;
22+
import org.testcontainers.utility.DockerImageName;
23+
24+
import static org.assertj.core.api.Assertions.assertThat;
25+
26+
class AwsOpenSearchDockerComposeConnectionDetailsFactoryTests extends AbstractDockerComposeIntegrationTests {
27+
28+
AwsOpenSearchDockerComposeConnectionDetailsFactoryTests() {
29+
super("localstack-compose.yaml", DockerImageName.parse("localstack/localstack:3.5.0"));
30+
}
31+
32+
@Test
33+
void runCreatesConnectionDetails() {
34+
AwsOpenSearchConnectionDetails connectionDetails = run(AwsOpenSearchConnectionDetails.class);
35+
assertThat(connectionDetails.getAccessKey()).isEqualTo("test");
36+
assertThat(connectionDetails.getSecretKey()).isEqualTo("test");
37+
assertThat(connectionDetails.getRegion()).isEqualTo("us-east-1");
38+
}
39+
40+
}

0 commit comments

Comments
 (0)