From 6cb89d74525b00cd00a368b0d7f3b38d4b862ea8 Mon Sep 17 00:00:00 2001 From: Christoph Strobl Date: Fri, 8 May 2020 09:30:25 +0200 Subject: [PATCH] DATAMONGO-2536 - Add option to skip reading aggregation result. Introduce dedicated AggregationPipeline to encapsulate pipeline stages. Original pull request: #876. --- .../data/mongodb/core/MongoTemplate.java | 12 ++ .../mongodb/core/ReactiveMongoTemplate.java | 13 +- .../mongodb/core/aggregation/Aggregation.java | 28 ++--- .../aggregation/AggregationOperation.java | 4 + .../core/aggregation/AggregationOptions.java | 42 +++++++ .../core/aggregation/AggregationPipeline.java | 111 ++++++++++++++++++ .../core/aggregation/AggregationUpdate.java | 6 +- .../core/aggregation/BucketAutoOperation.java | 7 +- .../core/aggregation/BucketOperation.java | 7 +- .../core/aggregation/CountOperation.java | 7 +- .../DocumentEnhancingOperation.java | 5 + .../core/aggregation/FacetOperation.java | 7 +- .../core/aggregation/GeoNearOperation.java | 7 +- .../aggregation/GraphLookupOperation.java | 7 +- .../core/aggregation/GroupOperation.java | 7 +- .../core/aggregation/LimitOperation.java | 7 +- .../core/aggregation/LookupOperation.java | 7 +- .../core/aggregation/MatchOperation.java | 7 +- .../core/aggregation/MergeOperation.java | 9 +- .../core/aggregation/OutOperation.java | 7 +- .../core/aggregation/ProjectionOperation.java | 7 +- .../core/aggregation/RedactOperation.java | 7 +- .../aggregation/ReplaceRootOperation.java | 5 + .../core/aggregation/SampleOperation.java | 7 +- .../core/aggregation/SkipOperation.java | 7 +- .../aggregation/SortByCountOperation.java | 7 +- .../core/aggregation/SortOperation.java | 7 +- .../core/aggregation/TypedAggregation.java | 2 +- .../core/aggregation/UnsetOperation.java | 9 +- .../core/aggregation/UnwindOperation.java | 9 +- .../core/aggregation/AggregationTests.java | 25 +++- .../aggregation/ReactiveAggregationTests.java | 18 +++ 32 files changed, 366 insertions(+), 51 deletions(-) create mode 100644 spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationPipeline.java 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 015cbf25b..1b1ad7d33 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 @@ -62,6 +62,7 @@ import org.springframework.data.mongodb.core.QueryOperations.UpdateContext; import org.springframework.data.mongodb.core.aggregation.Aggregation; import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext; import org.springframework.data.mongodb.core.aggregation.AggregationOptions; +import org.springframework.data.mongodb.core.aggregation.AggregationOptions.ResultOptions; import org.springframework.data.mongodb.core.aggregation.AggregationResults; import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext; import org.springframework.data.mongodb.core.aggregation.TypedAggregation; @@ -2152,6 +2153,17 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware, aggregateIterable = aggregateIterable.maxTime(options.getMaxTime().toMillis(), TimeUnit.MILLISECONDS); } + if(ResultOptions.SKIP.equals(options.resultOptions())) { + + // toCollection only allowed for $out and $merge if those are the last stages + if(aggregation.getPipeline().isOutOrMerge()) { + aggregateIterable.toCollection(); + } else { + aggregateIterable.first(); + } + return new AggregationResults<>(Collections.emptyList(), new Document()); + } + MongoIterable iterable = aggregateIterable.map(val -> { rawResult.add(val); 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 03daa82cb..4688df6c7 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 @@ -26,6 +26,7 @@ import java.util.*; import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import java.util.function.Function; +import java.util.function.Supplier; import java.util.stream.Collectors; import org.bson.BsonValue; @@ -68,6 +69,7 @@ import org.springframework.data.mongodb.core.QueryOperations.UpdateContext; import org.springframework.data.mongodb.core.aggregation.Aggregation; import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext; import org.springframework.data.mongodb.core.aggregation.AggregationOptions; +import org.springframework.data.mongodb.core.aggregation.AggregationOptions.ResultOptions; import org.springframework.data.mongodb.core.aggregation.PrefixingDelegatingAggregationOperationContext; import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext; import org.springframework.data.mongodb.core.aggregation.TypedAggregation; @@ -1004,11 +1006,11 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati } ReadDocumentCallback readCallback = new ReadDocumentCallback<>(mongoConverter, outputType, collectionName); - return execute(collectionName, collection -> aggregateAndMap(collection, pipeline, options, readCallback, + return execute(collectionName, collection -> aggregateAndMap(collection, pipeline, () -> aggregation.getPipeline().isOutOrMerge(), options, readCallback, aggregation instanceof TypedAggregation ? ((TypedAggregation) aggregation).getInputType() : null)); } - private Flux aggregateAndMap(MongoCollection collection, List pipeline, + private Flux aggregateAndMap(MongoCollection collection, List pipeline, Supplier isOutOrMerge, AggregationOptions options, ReadDocumentCallback readCallback, @Nullable Class inputType) { AggregatePublisher cursor = collection.aggregate(pipeline, Document.class) @@ -1028,6 +1030,13 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati cursor = cursor.maxTime(options.getMaxTime().toMillis(), TimeUnit.MILLISECONDS); } + if (ResultOptions.SKIP.equals(options.resultOptions())) { + if (isOutOrMerge.get()) { + return Flux.from(cursor.toCollection()).map(it -> (O) it); + } + return Flux.from(cursor.first()).thenMany(Mono.empty()); + } + return Flux.from(cursor).concatMap(readCallback::doWith); } diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Aggregation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Aggregation.java index 912897741..32c4a441c 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Aggregation.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Aggregation.java @@ -96,7 +96,7 @@ public class Aggregation { public static final AggregationOperationContext DEFAULT_CONTEXT = AggregationOperationRenderer.DEFAULT_CONTEXT; public static final AggregationOptions DEFAULT_OPTIONS = newAggregationOptions().build(); - protected final List operations; + protected final AggregationPipeline pipeline; private final AggregationOptions options; /** @@ -139,7 +139,7 @@ public class Aggregation { public Aggregation withOptions(AggregationOptions options) { Assert.notNull(options, "AggregationOptions must not be null."); - return new Aggregation(this.operations, options); + return new Aggregation(this.pipeline.getOperations(), options); } /** @@ -202,26 +202,10 @@ public class Aggregation { Assert.notNull(aggregationOperations, "AggregationOperations must not be null!"); Assert.notNull(options, "AggregationOptions must not be null!"); - // check $out/$merge is the last operation if it exists - for (AggregationOperation aggregationOperation : aggregationOperations) { - - if (aggregationOperation instanceof OutOperation && !isLast(aggregationOperation, aggregationOperations)) { - throw new IllegalArgumentException("The $out operator must be the last stage in the pipeline."); - } - - if (aggregationOperation instanceof MergeOperation && !isLast(aggregationOperation, aggregationOperations)) { - throw new IllegalArgumentException("The $merge operator must be the last stage in the pipeline."); - } - } - - this.operations = aggregationOperations; + this.pipeline = new AggregationPipeline(aggregationOperations); this.options = options; } - private boolean isLast(AggregationOperation aggregationOperation, List aggregationOperations) { - return aggregationOperations.indexOf(aggregationOperation) == aggregationOperations.size() - 1; - } - /** * Get the {@link AggregationOptions}. * @@ -718,7 +702,11 @@ public class Aggregation { * @since 2.1 */ public List toPipeline(AggregationOperationContext rootContext) { - return AggregationOperationRenderer.toDocument(operations, rootContext); + return pipeline.toDocuments(rootContext); + } + + public AggregationPipeline getPipeline() { + return pipeline; } /** diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOperation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOperation.java index 2b8a7171e..e611528e7 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOperation.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOperation.java @@ -54,4 +54,8 @@ public interface AggregationOperation { default List toPipelineStages(AggregationOperationContext context) { return Collections.singletonList(toDocument(context)); } + + default String operator() { + return toDocument(Aggregation.DEFAULT_CONTEXT).keySet().iterator().next(); + } } 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 96b199b8c..aab74a8a6 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 @@ -52,6 +52,7 @@ public class AggregationOptions { private final Optional collation; private final Optional comment; private Duration maxTime = Duration.ZERO; + private ResultOptions resultOptions = ResultOptions.READ; /** * Creates a new {@link AggregationOptions}. @@ -219,6 +220,14 @@ public class AggregationOptions { return maxTime; } + /** + * @return the {@link ResultOptions} to be used when running the {@link Aggregation}. Never {@literal null}. + * @since 3.0 + */ + public ResultOptions resultOptions() { + return resultOptions; + } + /** * Returns a new potentially adjusted copy for the given {@code aggregationCommandObject} with the configuration * applied. @@ -309,6 +318,7 @@ public class AggregationOptions { private @Nullable Collation collation; private @Nullable String comment; private @Nullable Duration maxTime; + private @Nullable ResultOptions resultOptions; /** * Defines whether to off-load intensive sort-operations to disk. @@ -399,6 +409,20 @@ public class AggregationOptions { return this; } + /** + * Run the aggregation, but do NOT read the aggregation result from the store.
+ * If the expected result of the aggregation is rather large, eg. when using an {@literal $out} operation, this + * option allows to execute the aggregation without having the cursor return the operation result. + * + * @return this. + * @since 3.0 + */ + public Builder skipOutput() { + + this.resultOptions = ResultOptions.SKIP; + return this; + } + /** * Returns a new {@link AggregationOptions} instance with the given configuration. * @@ -410,8 +434,26 @@ public class AggregationOptions { if (maxTime != null) { options.maxTime = maxTime; } + if (resultOptions != null) { + options.resultOptions = resultOptions; + } return options; } } + + /** + * @since 3.0 + */ + public enum ResultOptions { + + /** + * Just do it!, and do not read the operation result. + */ + SKIP, + /** + * Read the aggregation result from the cursor. + */ + READ; + } } diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationPipeline.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationPipeline.java new file mode 100644 index 000000000..2b7ae7d2d --- /dev/null +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationPipeline.java @@ -0,0 +1,111 @@ +/* + * Copyright 2020 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.data.mongodb.core.aggregation; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import org.bson.Document; +import org.springframework.util.Assert; + +/** + * The {@link AggregationPipeline} holds the collection of {@link AggregationOperation aggregation stages}. + * + * @author Christoph Strobl + * @since 3.1 + */ +public class AggregationPipeline { + + private final List pipeline; + + /** + * Create an empty pipeline + */ + public AggregationPipeline() { + this(new ArrayList<>()); + } + + /** + * Create a new pipeline with given {@link AggregationOperation stages}. + * + * @param aggregationOperations must not be {@literal null}. + */ + public AggregationPipeline(List aggregationOperations) { + pipeline = new ArrayList<>(aggregationOperations); + } + + /** + * Append the given {@link AggregationOperation stage} to the pipeline. + * + * @param aggregationOperation must not be {@literal null}. + * @return this. + */ + public AggregationPipeline add(AggregationOperation aggregationOperation) { + + Assert.notNull(aggregationOperation, "AggregationOperation must not be null!"); + + pipeline.add(aggregationOperation); + return this; + } + + /** + * Get the list of {@link AggregationOperation aggregation stages}. + * + * @return never {@literal null}. + */ + public List getOperations() { + return Collections.unmodifiableList(pipeline); + } + + List toDocuments(AggregationOperationContext context) { + + verify(); + return AggregationOperationRenderer.toDocument(pipeline, context); + } + + /** + * @return {@literal true} if the last aggregation stage is either {@literal $out} or {@literal $merge}. + */ + public boolean isOutOrMerge() { + + if (pipeline.isEmpty()) { + return false; + } + + String operator = pipeline.get(pipeline.size() - 1).operator(); + return operator.equals("$out") || operator.equals("$merge"); + } + + void verify() { + + // check $out/$merge is the last operation if it exists + for (AggregationOperation aggregationOperation : pipeline) { + + if (aggregationOperation instanceof OutOperation && !isLast(aggregationOperation)) { + throw new IllegalArgumentException("The $out operator must be the last stage in the pipeline."); + } + + if (aggregationOperation instanceof MergeOperation && !isLast(aggregationOperation)) { + throw new IllegalArgumentException("The $merge operator must be the last stage in the pipeline."); + } + } + } + + private boolean isLast(AggregationOperation aggregationOperation) { + return pipeline.indexOf(aggregationOperation) == pipeline.size() - 1; + } +} diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationUpdate.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationUpdate.java index e5dadc841..6efc28106 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationUpdate.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationUpdate.java @@ -139,7 +139,7 @@ public class AggregationUpdate extends Aggregation implements UpdateDefinition { setOperation.getFields().forEach(it -> { keysTouched.add(it.getName()); }); - operations.add(setOperation); + pipeline.add(setOperation); return this; } @@ -155,7 +155,7 @@ public class AggregationUpdate extends Aggregation implements UpdateDefinition { Assert.notNull(unsetOperation, "UnsetOperation must not be null!"); - operations.add(unsetOperation); + pipeline.add(unsetOperation); keysTouched.addAll(unsetOperation.removedFieldNames()); return this; } @@ -172,7 +172,7 @@ public class AggregationUpdate extends Aggregation implements UpdateDefinition { public AggregationUpdate replaceWith(ReplaceWithOperation replaceWithOperation) { Assert.notNull(replaceWithOperation, "ReplaceWithOperation must not be null!"); - operations.add(replaceWithOperation); + pipeline.add(replaceWithOperation); return this; } diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/BucketAutoOperation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/BucketAutoOperation.java index bb7e032ef..22511adb8 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/BucketAutoOperation.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/BucketAutoOperation.java @@ -106,7 +106,12 @@ public class BucketAutoOperation extends BucketOperationSupport extends Aggregation { public TypedAggregation withOptions(AggregationOptions options) { Assert.notNull(options, "AggregationOptions must not be null."); - return new TypedAggregation(inputType, operations, options); + return new TypedAggregation(inputType, pipeline.getOperations(), options); } } diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnsetOperation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnsetOperation.java index abd62c5b5..89dc79ea8 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnsetOperation.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnsetOperation.java @@ -117,13 +117,18 @@ public class UnsetOperation implements InheritsFieldsAggregationOperation { public Document toDocument(AggregationOperationContext context) { if (fields.size() == 1) { - return new Document("$unset", computeFieldName(fields.iterator().next(), context)); + return new Document(operator(), computeFieldName(fields.iterator().next(), context)); } - return new Document("$unset", + return new Document(operator(), fields.stream().map(it -> computeFieldName(it, context)).collect(Collectors.toList())); } + @Override + public String operator() { + return "$unset"; + } + private Object computeFieldName(Object field, AggregationOperationContext context) { if (field instanceof Field) { diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnwindOperation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnwindOperation.java index 7bb36f999..492afd244 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnwindOperation.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnwindOperation.java @@ -94,7 +94,7 @@ public class UnwindOperation String path = context.getReference(field).toString(); if (!preserveNullAndEmptyArrays && arrayIndex == null) { - return new Document("$unwind", path); + return new Document(operator(), path); } Document unwindArgs = new Document(); @@ -104,7 +104,12 @@ public class UnwindOperation } unwindArgs.put("preserveNullAndEmptyArrays", preserveNullAndEmptyArrays); - return new Document("$unwind", unwindArgs); + return new Document(operator(), unwindArgs); + } + + @Override + public String operator() { + return "$unwind"; } /* diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java index 33e9fdb34..6be58d731 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java @@ -1625,12 +1625,12 @@ public class AggregationTests { mongoTemplate.save(new Person("Leoniv", "Yakubov", 55, Person.Sex.MALE)); } - @Test // DATAMONGO-1418 + @Test // DATAMONGO-1418, DATAMONGO-2536 public void outShouldOutBeTheLastOperation() { assertThatIllegalArgumentException().isThrownBy(() -> newAggregation(match(new Criteria()), // group("field1").count().as("totalCount"), // out("collection1"), // - skip(100L))); + skip(100L)).toPipeline(DEFAULT_CONTEXT)); } @Test // DATAMONGO-1325 @@ -1907,6 +1907,27 @@ public class AggregationTests { assertThat(result.getMappedResults()).containsOnly(source); } + @Test // DATAMONGO-2536 + public void skipOutputDoesNotReadBackAggregationResults() { + + createTagDocuments(); + + Aggregation agg = newAggregation( // + project("tags"), // + unwind("tags"), // + group("tags") // + .count().as("n"), // + project("n") // + .and("tag").previousOperation(), // + sort(DESC, "n") // + ).withOptions(AggregationOptions.builder().skipOutput().build()); + + AggregationResults results = mongoTemplate.aggregate(agg, INPUT_COLLECTION, TagCount.class); + + assertThat(results.getMappedResults()).isEmpty(); + assertThat(results.getRawResults()).isEmpty(); + } + private void createUsersWithReferencedPersons() { mongoTemplate.dropCollection(User.class); diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/ReactiveAggregationTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/ReactiveAggregationTests.java index 262750e63..61d3834b4 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/ReactiveAggregationTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/ReactiveAggregationTests.java @@ -164,4 +164,22 @@ public class ReactiveAggregationTests { reactiveMongoTemplate.aggregate(aggregation, "newyork", Document.class).as(StepVerifier::create).expectNextCount(4) .verifyComplete(); } + + @Test // DATAMONGO-2356 + public void skipOutputDoesNotReadBackAggregationResults() { + + Product product = new Product("P1", "A", 1.99, 3, 0.05, 0.19); + reactiveMongoTemplate.insert(product).as(StepVerifier::create).expectNextCount(1).verifyComplete(); + + double shippingCosts = 1.2; + + TypedAggregation agg = newAggregation(Product.class, // + project("name", "netPrice") // + .andExpression("netPrice * 10", shippingCosts).as("salesPrice") // + ).withOptions(AggregationOptions.builder().skipOutput().build()); + + reactiveMongoTemplate.aggregate(agg, Document.class).as(StepVerifier::create).verifyComplete(); + } + + }