From 4220df5bf8cb906671ef144b74141cb2b20572b4 Mon Sep 17 00:00:00 2001 From: Christoph Strobl Date: Wed, 30 Nov 2022 15:54:29 +0100 Subject: [PATCH] Accept index names as hint for aggregations. Closes #4238 Original pull request: #4243 --- .../data/mongodb/core/MongoTemplate.java | 32 ++++++++-- .../mongodb/core/ReactiveMongoTemplate.java | 11 +++- .../core/aggregation/AggregationOptions.java | 58 ++++++++++++++++++- .../mongodb/core/MongoTemplateUnitTests.java | 12 ++++ .../core/ReactiveMongoTemplateUnitTests.java | 13 +++++ 5 files changed, 116 insertions(+), 10 deletions(-) diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java index 5d642201a..9930e6243 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java @@ -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 createView(String name, Class source, AggregationPipeline pipeline, @Nullable ViewOptions options) { + public MongoCollection 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 createView(String name, String source, AggregationPipeline pipeline, @Nullable ViewOptions options) { + public MongoCollection 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 doCreateView(String name, String source, List pipeline, @Nullable ViewOptions options) { + protected MongoCollection doCreateView(String name, String source, List 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; diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java index 1f7c015fd..e69f51754 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java @@ -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) // diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOptions.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOptions.java index cf3a82438..6ea2743f9 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOptions.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOptions.java @@ -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 cursor; private final Optional collation; private final Optional comment; - private final Optional hint; + private final Optional 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 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 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 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. * diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java index 770819ac0..850797880 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java @@ -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() { diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java index 2f9ce1ac8..df1bb8446 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/ReactiveMongoTemplateUnitTests.java @@ -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() {