Spring/spring-data-mongodb
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Accept index names as hint for aggregations.

Browse Source 
Closes #4238
Original pull request: #4243
This commit is contained in:
Christoph Strobl
2022-11-30 15:54:29 +01:00
committed by Mark Paluch
parent 95c6d1531f
commit 4220df5bf8
5 changed files with 116 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java
View File

@@ -30,7 +30,6 @@ import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.bson.Document;
import org.bson.conversions.Bson;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
@@ -634,7 +633,8 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
}
@Override
public MongoCollection<Document> createView(String name, Class<?> source, AggregationPipeline pipeline, @Nullable ViewOptions options) {
public MongoCollection<Document> createView(String name, Class<?> source, AggregationPipeline pipeline,
@Nullable ViewOptions options) {
return createView(name, getCollectionName(source),
queryOperations.createAggregation(Aggregation.newAggregation(source, pipeline.getOperations()), source),
@@ -642,7 +642,8 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
}
@Override
public MongoCollection<Document> createView(String name, String source, AggregationPipeline pipeline, @Nullable ViewOptions options) {
public MongoCollection<Document> createView(String name, String source, AggregationPipeline pipeline,
@Nullable ViewOptions options) {
return createView(name, source,
queryOperations.createAggregation(Aggregation.newAggregation(pipeline.getOperations()), (Class<?>) null),
@@ -654,7 +655,8 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
return doCreateView(name, source, aggregation.getAggregationPipeline(), options);
}
protected MongoCollection<Document> doCreateView(String name, String source, List<Document> pipeline, @Nullable ViewOptions options) {
protected MongoCollection<Document> doCreateView(String name, String source, List<Document> pipeline,
@Nullable ViewOptions options) {
CreateViewOptions viewOptions = new CreateViewOptions();
if (options != null) {
@@ -2065,7 +2067,16 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
}
options.getComment().ifPresent(aggregateIterable::comment);
options.getHint().ifPresent(aggregateIterable::hint);
if (options.getHintObject().isPresent()) {
Object hintObject = options.getHintObject().get();
if (hintObject instanceof String hintString) {
aggregateIterable = aggregateIterable.hintString(hintString);
} else if (hintObject instanceof Document hintDocument) {
aggregateIterable = aggregateIterable.hint(hintDocument);
} else {
throw new IllegalStateException("Unable to read hint of type %s".formatted(hintObject.getClass()));
}
}
if (options.hasExecutionTimeLimit()) {
aggregateIterable = aggregateIterable.maxTime(options.getMaxTime().toMillis(), TimeUnit.MILLISECONDS);
@@ -2124,7 +2135,16 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
}
options.getComment().ifPresent(cursor::comment);
options.getHint().ifPresent(cursor::hint);
if (options.getHintObject().isPresent()) {
Object hintObject = options.getHintObject().get();
if (hintObject instanceof String hintString) {
cursor = cursor.hintString(hintString);
} else if (hintObject instanceof Document hintDocument) {
cursor = cursor.hint(hintDocument);
} else {
throw new IllegalStateException("Unable to read hint of type %s".formatted(hintObject.getClass()));
}
}
Class<?> domainType = aggregation instanceof TypedAggregation ? ((TypedAggregation) aggregation).getInputType()
: null;

11
spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java
View File

@@ -938,7 +938,16 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati
}
options.getComment().ifPresent(cursor::comment);
options.getHint().ifPresent(cursor::hint);
if (options.getHintObject().isPresent()) {
Object hintObject = options.getHintObject().get();
if (hintObject instanceof String hintString) {
cursor = cursor.hintString(hintString);
} else if (hintObject instanceof Document hintDocument) {
cursor = cursor.hint(hintDocument);
} else {
throw new IllegalStateException("Unable to read hint of type %s".formatted(hintObject.getClass()));
}
}
Optionals.firstNonEmpty(options::getCollation, () -> operations.forType(inputType).getCollation()) //
.map(Collation::toMongoCollation) //

58
spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOptions.java
View File

@@ -20,6 +20,7 @@ import java.util.Optional;
import org.bson.Document;
import org.springframework.data.mongodb.core.query.Collation;
import org.springframework.data.mongodb.util.BsonUtils;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
@@ -53,7 +54,7 @@ public class AggregationOptions {
private final Optional<Document> cursor;
private final Optional<Collation> collation;
private final Optional<String> comment;
private final Optional<Document> hint;
private final Optional<Object> hint;
private Duration maxTime = Duration.ZERO;
private ResultOptions resultOptions = ResultOptions.READ;
private DomainTypeMapping domainTypeMapping = DomainTypeMapping.RELAXED;
@@ -113,7 +114,7 @@ public class AggregationOptions {
* @since 3.1
*/
private AggregationOptions(boolean allowDiskUse, boolean explain, @Nullable Document cursor,
@Nullable Collation collation, @Nullable String comment, @Nullable Document hint) {
@Nullable Collation collation, @Nullable String comment, @Nullable Object hint) {
this.allowDiskUse = allowDiskUse;
this.explain = explain;
@@ -242,6 +243,44 @@ public class AggregationOptions {
* @since 3.1
*/
public Optional<Document> getHint() {
return hint.map(it -> {
if (it instanceof Document doc) {
return doc;
}
if (it instanceof String hintString) {
if (BsonUtils.isJsonDocument(hintString)) {
return BsonUtils.parse(hintString, null);
}
}
throw new IllegalStateException("Unable to read hint of type %s".formatted(it.getClass()));
});
}
/**
* Get the hint (indexName) used to to fulfill the aggregation.
*
* @return never {@literal null}.
* @since 4.1
*/
public Optional<String> getHintAsString() {
return hint.map(it -> {
if (it instanceof String hintString) {
return hintString;
}
if (it instanceof Document doc) {
return BsonUtils.toJson(doc);
}
throw new IllegalStateException("Unable to read hint of type %s".formatted(it.getClass()));
});
}
/**
* Get the hint used to to fulfill the aggregation.
*
* @return never {@literal null}.
* @since 4.1
*/
public Optional<Object> getHintObject() {
return hint;
}
@@ -361,7 +400,7 @@ public class AggregationOptions {
private @Nullable Document cursor;
private @Nullable Collation collation;
private @Nullable String comment;
private @Nullable Document hint;
private @Nullable Object hint;
private @Nullable Duration maxTime;
private @Nullable ResultOptions resultOptions;
private @Nullable DomainTypeMapping domainTypeMapping;
@@ -454,6 +493,19 @@ public class AggregationOptions {
return this;
}
/**
* Define a hint that is used by query optimizer to to fulfill the aggregation.
*
* @param indexName can be {@literal null}.
* @return this.
* @since 4.1
*/
public Builder hint(@Nullable String indexName) {
this.hint = indexName;
return this;
}
/**
* Set the time limit for processing.
*

12
spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java
View File

@@ -204,6 +204,8 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
when(aggregateIterable.map(any())).thenReturn(aggregateIterable);
when(aggregateIterable.maxTime(anyLong(), any())).thenReturn(aggregateIterable);
when(aggregateIterable.into(any())).thenReturn(Collections.emptyList());
when(aggregateIterable.hint(any())).thenReturn(aggregateIterable);
when(aggregateIterable.hintString(any())).thenReturn(aggregateIterable);
when(distinctIterable.collation(any())).thenReturn(distinctIterable);
when(distinctIterable.map(any())).thenReturn(distinctIterable);
when(distinctIterable.into(any())).thenReturn(Collections.emptyList());
@@ -497,6 +499,16 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
verify(aggregateIterable).hint(hint);
}
@Test // GH-4238
void aggregateShouldHonorOptionsHintString() {
AggregationOptions options = AggregationOptions.builder().hint("index-1").build();
template.aggregate(newAggregation(Aggregation.unwind("foo")).withOptions(options), "collection-1", Wrapper.class);
verify(aggregateIterable).hintString("index-1");
}
@Test // GH-3542
void aggregateShouldUseRelaxedMappingByDefault() {

13
spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java
View File

@@ -23,6 +23,8 @@ import static org.springframework.data.mongodb.test.util.Assertions.assertThat;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.springframework.data.mongodb.core.MongoTemplateUnitTests.Wrapper;
import org.springframework.data.mongodb.core.aggregation.Aggregation;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;
@@ -666,6 +668,17 @@ public class ReactiveMongoTemplateUnitTests {
verify(aggregatePublisher).hint(hint);
}
@Test // GH-4238
void aggregateShouldHonorOptionsHintString() {
AggregationOptions options = AggregationOptions.builder().hint("index-1").build();
template.aggregate(newAggregation(Sith.class, project("id")).withOptions(options), AutogenerateableId.class,
Document.class).subscribe();
verify(aggregatePublisher).hintString("index-1");
}
@Test // DATAMONGO-2390
void aggregateShouldNoApplyZeroOrNegativeMaxTime() {