Skip to content

Commit b74f1d7

Browse files
christophstroblchristophstrobl
committed
Prepare for integration tests
1 parent 9b745c4 commit b74f1d7

File tree

11 files changed

+295
-76
lines changed

11 files changed

+295
-76
lines changed

‎spring-data-mongodb/pom.xml

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,20 @@
273273
<scope>test</scope>
274274
</dependency>
275275

276+
<dependency>
277+
<groupId>org.testcontainers</groupId>
278+
<artifactId>junit-jupiter</artifactId>
279+
<version>${testcontainers}</version>
280+
<scope>test</scope>
281+
</dependency>
282+
283+
<dependency>
284+
<groupId>org.testcontainers</groupId>
285+
<artifactId>mongodb</artifactId>
286+
<version>${testcontainers}</version>
287+
<scope>test</scope>
288+
</dependency>
289+
276290
<dependency>
277291
<groupId>jakarta.transaction</groupId>
278292
<artifactId>jakarta.transaction-api</artifactId>

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/index/DefaultSearchIndexOperations.java

Lines changed: 35 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2024 the original author or authors.
2+
* Copyright 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.
@@ -18,24 +18,23 @@
1818
import java.util.ArrayList;
1919
import java.util.List;
2020

21+
import org.bson.BsonString;
2122
import org.bson.Document;
22-
2323
import org.springframework.data.mapping.context.MappingContext;
2424
import org.springframework.data.mongodb.core.MongoOperations;
25-
import org.springframework.data.mongodb.core.aggregation.Aggregation;
26-
import org.springframework.data.mongodb.core.aggregation.AggregationResults;
2725
import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
2826
import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
2927
import org.springframework.data.util.TypeInformation;
3028
import org.springframework.lang.Nullable;
29+
import org.springframework.util.StringUtils;
3130

3231
import com.mongodb.client.model.SearchIndexModel;
3332
import com.mongodb.client.model.SearchIndexType;
3433

3534
/**
3635
* @author Christoph Strobl
3736
* @author Mark Paluch
38-
* @since 3.5
37+
* @since 4.5
3938
*/
4039
public class DefaultSearchIndexOperations implements SearchIndexOperations {
4140

@@ -65,17 +64,14 @@ public DefaultSearchIndexOperations(MongoOperations mongoOperations, String coll
6564
@Override
6665
public String ensureIndex(SearchIndexDefinition indexDefinition) {
6766

68-
if (!(indexDefinition instanceof VectorIndex vsi)) {
69-
throw new IllegalStateException("Index definitions must be of type VectorIndex");
70-
}
71-
7267
Document index = indexDefinition.getIndexDocument(entityTypeInformation,
7368
mongoOperations.getConverter().getMappingContext());
7469

75-
mongoOperations.getCollection(collectionName).createSearchIndexes(List
76-
.of(new SearchIndexModel(vsi.getName(), (Document) index.get("definition"), SearchIndexType.vectorSearch())));
70+
mongoOperations.getCollection(collectionName)
71+
.createSearchIndexes(List.of(new SearchIndexModel(indexDefinition.getName(), index.get("definition", Document.class),
72+
SearchIndexType.of(new BsonString(indexDefinition.getType())))));
7773

78-
return vsi.getName();
74+
return indexDefinition.getName();
7975
}
8076

8177
@Override
@@ -94,7 +90,7 @@ public void updateIndex(SearchIndexDefinition index) {
9490
@Override
9591
public boolean exists(String indexName) {
9692

97-
List<Document> indexes = mongoOperations.getCollection(collectionName).listSearchIndexes().into(new ArrayList<>());
93+
List<Document> indexes = getSearchIndexes(indexName);
9894

9995
for (Document index : indexes) {
10096
if (index.getString("name").equals(indexName)) {
@@ -105,12 +101,27 @@ public boolean exists(String indexName) {
105101
return false;
106102
}
107103

104+
@Override
105+
public SearchIndexStatus status(String name) {
106+
107+
List<Document> indexes = getSearchIndexes(name);
108+
109+
if (indexes.isEmpty()) {
110+
return SearchIndexStatus.DOES_NOT_EXIST;
111+
}
112+
113+
for (Document index : indexes) {
114+
if (index.getString("name").equals(name)) {
115+
return SearchIndexStatus.valueOf(index.getString("status"));
116+
}
117+
}
118+
return SearchIndexStatus.DOES_NOT_EXIST;
119+
}
120+
108121
@Override
109122
public List<IndexInfo> getIndexInfo() {
110123

111-
AggregationResults<Document> aggregate = mongoOperations.aggregate(
112-
Aggregation.newAggregation(context -> new Document("$listSearchIndexes", new Document())), collectionName,
113-
Document.class);
124+
List<Document> aggregate = getSearchIndexes(null);
114125

115126
ArrayList<IndexInfo> result = new ArrayList<>();
116127
for (Document doc : aggregate) {
@@ -139,4 +150,12 @@ public void dropIndex(String name) {
139150
mongoOperations.getCollection(collectionName).dropSearchIndex(name);
140151
}
141152

153+
private List<Document> getSearchIndexes(String indexName) {
154+
155+
Document filter = StringUtils.hasText(indexName) ? new Document("name", indexName) : new Document();
156+
157+
return mongoOperations.getCollection(collectionName).aggregate(List.of(new Document("$listSearchIndexes", filter)))
158+
.into(new ArrayList<>());
159+
}
160+
142161
}

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/index/SearchIndexDefinition.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2011-2024 the original author or authors.
2+
* Copyright 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.
@@ -73,4 +73,6 @@ default Document getIndexDocument(@Nullable TypeInformation<?> entity,
7373
Document getDefinition(@Nullable TypeInformation<?> entity,
7474
MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext);
7575

76+
77+
7678
}

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/index/SearchIndexOperations.java

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
/*
2-
* Copyright 2024. the original author or authors.
2+
* Copyright 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.
66
* You may obtain a copy of the License at
77
*
8-
* http://www.apache.org/licenses/LICENSE-2.0
8+
* https://www.apache.org/licenses/LICENSE-2.0
99
*
1010
* Unless required by applicable law or agreed to in writing, software
1111
* distributed under the License is distributed on an "AS IS" BASIS,
@@ -53,6 +53,9 @@ public interface SearchIndexOperations {
5353
*/
5454
boolean exists(String name);
5555

56+
57+
SearchIndexStatus status(String name);
58+
5659
/**
5760
* Drops an index from this collection.
5861
*

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/index/SearchIndexOperationsProvider.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
/*
2-
* Copyright 2024. the original author or authors.
2+
* Copyright 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.
66
* You may obtain a copy of the License at
77
*
8-
* http://www.apache.org/licenses/LICENSE-2.0
8+
* https://www.apache.org/licenses/LICENSE-2.0
99
*
1010
* Unless required by applicable law or agreed to in writing, software
1111
* distributed under the License is distributed on an "AS IS" BASIS,

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/index/SearchIndexStatus.java

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/*
2+
* Copyright 2025. 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+
* http://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+
package org.springframework.data.mongodb.core.index;
17+
18+
/**
19+
* @author Christoph Strobl
20+
* @since 2025/01
21+
*/
22+
public enum SearchIndexStatus {
23+
24+
/** building or re-building the index - might be queryable */
25+
BUILDING,
26+
27+
/** nothing to be seen here - not queryable */
28+
DOES_NOT_EXIST,
29+
30+
/** will cease to exist - no longer queryable */
31+
DELETING,
32+
33+
/** well, that did not work - not queryable */
34+
FAILED,
35+
36+
/** busy with other things, check back later - not queryable */
37+
PENDING,
38+
39+
/** ask me anything - queryable */
40+
READY,
41+
42+
/** ask me anything about outdated data - still queryable */
43+
STALE
44+
}

‎spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/index/VectorIndex.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@
4545
import org.springframework.lang.Contract;
4646
import org.springframework.lang.Nullable;
4747
import org.springframework.util.Assert;
48+
import org.springframework.util.ObjectUtils;
49+
import org.springframework.util.StringUtils;
4850

4951
/**
5052
* {@link IndexDefinition} for creating MongoDB
@@ -144,7 +146,9 @@ public Document getDefinition(@Nullable TypeInformation<?> entity,
144146
filter.put("path", resolvePath(vif.path(), persistentEntity, mappingContext));
145147
filter.put("numDimensions", vif.dimensions());
146148
filter.put("similarity", vif.similarity());
147-
filter.put("quantization", vif.quantization());
149+
if(!ObjectUtils.nullSafeEquals(Quantization.NONE.getQuantizationName(), vif.quantization)) {
150+
filter.put("quantization", vif.quantization());
151+
}
148152
fields.add(filter);
149153
}
150154

‎spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/VectorSearchTests.java

Lines changed: 67 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -15,69 +15,27 @@
1515
*/
1616
package org.springframework.data.mongodb.core.aggregation;
1717

18+
import java.time.Duration;
19+
import java.util.List;
20+
1821
import org.bson.Document;
22+
import org.junit.jupiter.api.BeforeEach;
1923
import org.junit.jupiter.api.Test;
20-
import org.junit.jupiter.api.extension.ExtendWith;
21-
import org.springframework.data.mongodb.core.index.VectorIndex;
22-
import org.springframework.data.mongodb.test.util.EnableIfVectorSearchAvailable;
23-
import org.springframework.data.mongodb.test.util.MongoTemplateExtension;
24+
import org.springframework.data.mongodb.test.util.AtlasContainer;
2425
import org.springframework.data.mongodb.test.util.MongoTestTemplate;
25-
import org.springframework.data.mongodb.test.util.Template;
26+
import org.testcontainers.junit.jupiter.Container;
27+
import org.testcontainers.junit.jupiter.Testcontainers;
28+
29+
import com.mongodb.client.MongoClient;
30+
import com.mongodb.client.MongoClients;
2631

2732
/**
2833
* @author Christoph Strobl
2934
*/
30-
@EnableIfVectorSearchAvailable
31-
@ExtendWith(MongoTemplateExtension.class)
35+
@Testcontainers
3236
public class VectorSearchTests {
3337

3438
static final String COLLECTION_NAME = "movies";
35-
36-
@Template(database = "mflix") //
37-
static MongoTestTemplate template;
38-
39-
@Test
40-
void xxx() {
41-
42-
// boolean hasIndex = template.indexOps(COLLECTION_NAME).getIndexInfo().stream()
43-
// .anyMatch(it -> it.getName().endsWith("movie_vector_index"));
44-
45-
// TODO: index conversion etc. is missing - should we combine the index info listing?
46-
// boolean hasIndex = template.execute(db -> {
47-
//
48-
// Document doc = db.runCommand(new Document("listSearchIndexes", COLLECTION_NAME));
49-
// Object searchIndexes = BsonUtils.resolveValue(BsonUtils.asMap(doc), "cursor.firstBatch");
50-
// if(searchIndexes instanceof Collection<?> indexes) {
51-
// return indexes.stream().anyMatch(it -> it instanceof Document idx && idx.get("name",
52-
// String.class).equalsIgnoreCase("vector_index"));
53-
// }
54-
// return false;
55-
// });
56-
57-
if (!template.collectionExists(COLLECTION_NAME)) {
58-
template.createCollection(COLLECTION_NAME);
59-
}
60-
61-
boolean hasIndex = template.searchIndexOps(COLLECTION_NAME).exists("movie_vector_index");
62-
63-
if (!hasIndex) {
64-
65-
System.out.print("Creating index: ");
66-
VectorIndex vectorIndex = new VectorIndex("movie_vector_index").addVector("plot_embedding",
67-
field -> field.dimensions(1536).similarity(VectorIndex.SimilarityFunction.COSINE)).addFilter("language");
68-
String s = template.searchIndexOps(COLLECTION_NAME).ensureIndex(vectorIndex);
69-
}
70-
71-
VectorSearchOperation $vectorSearch = VectorSearchOperation.search("movie_vector_index").path("plot_embedding")
72-
.vector(vectors).limit(10).numCandidates(150).withSearchScore();
73-
74-
Aggregation agg = Aggregation.newAggregation($vectorSearch, Aggregation.project("plot", "title"));
75-
76-
AggregationResults<Document> aggregate = template.aggregate(agg, COLLECTION_NAME, Document.class);
77-
78-
aggregate.forEach(System.out::println);
79-
}
80-
8139
static double[] vectors = { -0.0016261312, -0.028070757, -0.011342932, -0.012775794, -0.0027440966, 0.008683807,
8240
-0.02575152, -0.02020668, -0.010283281, -0.0041719596, 0.021392956, 0.028657231, -0.006634482, 0.007490867,
8341
0.018593878, 0.0038187427, 0.029590257, -0.01451522, 0.016061379, 0.00008528442, -0.008943722, 0.01627464,
@@ -271,5 +229,61 @@ void xxx() {
271229
-0.0077707744, -0.012049366, 0.011869425, 0.00858384, -0.024698535, -0.030283362, 0.020140035, 0.011949399,
272230
-0.013968734, 0.042732596, -0.011649498, -0.011982721, -0.016967745, -0.0060913274, -0.007130985, -0.013109017,
273231
-0.009710136 };
232+
private static @Container AtlasContainer atlasLocal = new AtlasContainer();
233+
MongoClient client;
234+
MongoTestTemplate template;
235+
236+
@BeforeEach
237+
void beforeEach() {
238+
client = MongoClients.create(atlasLocal.getConnectionString());
239+
template = new MongoTestTemplate(client, "vector-search-tests");
240+
}
241+
242+
@Test
243+
void xxx() throws InterruptedException {
244+
245+
if (!template.collectionExists(COLLECTION_NAME)) {
246+
template.createCollection(COLLECTION_NAME);
247+
}
248+
249+
// TODO: index conversion etc. is missing - should we combine the index info listing?
250+
// boolean hasIndex = template.execute(db -> {
251+
//
252+
// Document doc = db.runCommand(new Document("listSearchIndexes", COLLECTION_NAME));
253+
// Object searchIndexes = BsonUtils.resolveValue(BsonUtils.asMap(doc), "cursor.firstBatch");
254+
// if (searchIndexes instanceof Collection<?> indexes) {
255+
// return indexes.stream().anyMatch(
256+
// it -> it instanceof Document idx && idx.get("name", String.class).equalsIgnoreCase("vector_index"));
257+
// }
258+
// return false;
259+
// });
260+
261+
boolean hasIndex = template.searchIndexOps(COLLECTION_NAME).exists("movie_vector_index");
262+
263+
if (!hasIndex) {
264+
265+
template.execute(db -> db.runCommand(createSearchIndexDefinition("movie_vector_index")).toJson());
266+
template.awaitIndexCreation(COLLECTION_NAME, "movie_vector_index", Duration.ofSeconds(5));
267+
}
268+
269+
VectorSearchOperation $vectorSearch = VectorSearchOperation.search("movie_vector_index").path("plot_embedding")
270+
.vector(vectors).limit(10).numCandidates(150).withSearchScore();
271+
272+
Aggregation agg = Aggregation.newAggregation($vectorSearch, Aggregation.project("plot", "title"));
273+
274+
AggregationResults<Document> aggregate = template.aggregate(agg, COLLECTION_NAME, Document.class);
275+
276+
aggregate.forEach(System.out::println);
277+
}
278+
279+
private org.bson.Document createSearchIndexDefinition(String indexName) {
280+
281+
List<Document> vectorFields = List.of(new org.bson.Document().append("type", "vector")
282+
.append("path", "plot_embedding").append("numDimensions", 1536).append("similarity", "cosine"));
283+
284+
return new org.bson.Document().append("createSearchIndexes", COLLECTION_NAME).append("indexes",
285+
List.of(new org.bson.Document().append("name", indexName).append("type", "vectorSearch").append("definition",
286+
new org.bson.Document("fields", vectorFields))));
287+
}
274288

275289
}

0 commit comments

Comments
 (0)