From 230c32041a38ad0fb142c645b603d8c8e0a893d2 Mon Sep 17 00:00:00 2001 From: Christoph Strobl Date: Thu, 17 Sep 2020 14:02:10 +0200 Subject: [PATCH] DATAMONGO-2622 - Add support for $unionWith aggregation stage. We now support the $unionWith aggregation stage via the UnionWithOperation that performs a union of two collections by combining pipeline results, potentially containing duplicates, into a single result set that is handed over to the next stage. In order to remove duplicates it is possible to append a GroupOperation right after UnionWithOperation. If the UnionWithOperation uses a pipeline to process documents, field names within the pipeline will be treated as is. In order to map domain type property names to actual field names (considering potential org.springframework.data.mongodb.core.mapping.Field annotations) make sure the enclosing aggregation is a TypedAggregation and provide the target type for the $unionWith stage via mapFieldsTo(Class). Original pull request: #886. --- .../data/mongodb/core/AggregationUtil.java | 14 +- .../data/mongodb/core/MongoTemplate.java | 4 +- .../core/aggregation/AggregationPipeline.java | 35 ++++ ...osedFieldsAggregationOperationContext.java | 8 + .../TypeBasedAggregationOperationContext.java | 13 ++ .../core/aggregation/UnionWithOperation.java | 170 ++++++++++++++++++ .../UnionWithOperationUnitTests.java | 129 +++++++++++++ 7 files changed, 366 insertions(+), 7 deletions(-) create mode 100644 spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnionWithOperation.java create mode 100644 spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/UnionWithOperationUnitTests.java diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/AggregationUtil.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/AggregationUtil.java index 0f956a317..b9c4be1d0 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/AggregationUtil.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/AggregationUtil.java @@ -28,6 +28,7 @@ import org.springframework.data.mongodb.core.aggregation.AggregationOperation; import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext; import org.springframework.data.mongodb.core.aggregation.AggregationOptions; import org.springframework.data.mongodb.core.aggregation.CountOperation; +import org.springframework.data.mongodb.core.aggregation.RelaxedTypeBasedAggregationOperationContext; import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext; import org.springframework.data.mongodb.core.aggregation.TypedAggregation; import org.springframework.data.mongodb.core.convert.QueryMapper; @@ -75,12 +76,17 @@ class AggregationUtil { return context; } - if (aggregation instanceof TypedAggregation) { - return new TypeBasedAggregationOperationContext(((TypedAggregation) aggregation).getInputType(), mappingContext, - queryMapper); + if (!(aggregation instanceof TypedAggregation)) { + return Aggregation.DEFAULT_CONTEXT; } - return Aggregation.DEFAULT_CONTEXT; + Class inputType = ((TypedAggregation) aggregation).getInputType(); + + if (aggregation.getPipeline().requiresRelaxedChecking()) { + return new RelaxedTypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper); + } + + return new TypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper); } /** 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 97deaae46..227ba0063 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 @@ -1977,9 +1977,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware, Assert.notNull(aggregation, "Aggregation pipeline must not be null!"); - AggregationOperationContext context = new TypeBasedAggregationOperationContext(aggregation.getInputType(), - mappingContext, queryMapper); - return aggregate(aggregation, inputCollectionName, outputType, context); + return aggregate(aggregation, inputCollectionName, outputType, null); } /* (non-Javadoc) 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 index 8c4b0ed10..083daa6d8 100644 --- 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 @@ -45,6 +45,8 @@ public class AggregationPipeline { * @param aggregationOperations must not be {@literal null}. */ public AggregationPipeline(List aggregationOperations) { + + Assert.notNull(aggregationOperations, "AggregationOperations must not be null!"); pipeline = new ArrayList<>(aggregationOperations); } @@ -108,4 +110,37 @@ public class AggregationPipeline { private boolean isLast(AggregationOperation aggregationOperation) { return pipeline.indexOf(aggregationOperation) == pipeline.size() - 1; } + + /** + * @return {@literal true} if field names might get computed by one of the pipeline stages, that the + * {@link AggregationOperationContext} might not be aware of. A strongly typed context might fail to resolve + * field references, so if {@literal true} usage of a {@link RelaxedTypeBasedAggregationOperationContext} + * might be the better choice. + * @since 3.1 + */ + public boolean requiresRelaxedChecking() { + return pipelineContainsValueOfType(UnionWithOperation.class); + } + + /** + * @return {@literal true} if the pipeline does not contain any stages. + * @since 3.1 + */ + public boolean isEmpty() { + return pipeline.isEmpty(); + } + + private boolean pipelineContainsValueOfType(Class type) { + + if (isEmpty()) { + return false; + } + + for (Object element : pipeline) { + if (type.isInstance(element)) { + return true; + } + } + return false; + } } diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/ExposedFieldsAggregationOperationContext.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/ExposedFieldsAggregationOperationContext.java index d802fa4d1..f0c2b60a2 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/ExposedFieldsAggregationOperationContext.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/ExposedFieldsAggregationOperationContext.java @@ -149,4 +149,12 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo } return null; } + + /** + * @return obtain the root context used to resolve references. + * @since 3.1 + */ + AggregationOperationContext getRootContext() { + return rootContext; + } } diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/TypeBasedAggregationOperationContext.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/TypeBasedAggregationOperationContext.java index 5ba22e960..aab2ffe5f 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/TypeBasedAggregationOperationContext.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/TypeBasedAggregationOperationContext.java @@ -133,6 +133,19 @@ public class TypeBasedAggregationOperationContext implements AggregationOperatio */ @Override public AggregationOperationContext continueOnMissingFieldReference() { + return continueOnMissingFieldReference(type); + } + + /** + * This toggle allows the {@link AggregationOperationContext context} to use any given field name without checking for + * its existence. Typically the {@link AggregationOperationContext} fails when referencing unknown fields, those that + * are not present in one of the previous stages or the input source, throughout the pipeline. + * + * @param type The domain type to map fields to. + * @return a more relaxed {@link AggregationOperationContext}. + * @since 3.1 + */ + public AggregationOperationContext continueOnMissingFieldReference(Class type) { return new RelaxedTypeBasedAggregationOperationContext(type, mappingContext, mapper); } diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnionWithOperation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnionWithOperation.java new file mode 100644 index 000000000..c30e8c7ec --- /dev/null +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/UnionWithOperation.java @@ -0,0 +1,170 @@ +/* + * 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.Arrays; +import java.util.List; + +import org.bson.Document; +import org.springframework.lang.Nullable; +import org.springframework.util.Assert; + +/** + * The $unionWith aggregation + * stage (available since MongoDB 4.4) performs a union of two collections by combining pipeline results, potentially + * containing duplicates, into a single result set that is handed over to the next stage.
+ * In order to remove duplicates it is possible to append a {@link GroupOperation} right after + * {@link UnionWithOperation}. + *

+ * If the {@link UnionWithOperation} uses a + * pipeline + * to process documents, field names within the pipeline will be treated as is. In order to map domain type property + * names to actual field names (considering potential {@link org.springframework.data.mongodb.core.mapping.Field} + * annotations) make sure the enclosing aggregation is a {@link TypedAggregation} and provide the target type for the + * {@code $unionWith} stage via {@link #mapFieldsTo(Class)}. + * + * @author Christoph Strobl + * @see Aggregation Pipeline Stage: + * $unionWith + * @since 3.1 + */ +public class UnionWithOperation implements AggregationOperation { + + private final String collection; + + @Nullable // + private final AggregationPipeline pipeline; + + @Nullable // + private final Class domainType; + + public UnionWithOperation(String collection, @Nullable AggregationPipeline pipeline, @Nullable Class domainType) { + + this.collection = collection; + this.pipeline = pipeline; + this.domainType = domainType; + } + + /** + * Set the name of the collection from which pipeline results should be included in the result set.
+ * The collection name is used to set the {@code coll} parameter of {@code $unionWith}. + * + * @param collection the MongoDB collection name. Must not be {@literal null}. + * @return new instance of {@link UnionWithOperation}. + * @throws IllegalArgumentException if the required argument is {@literal null}. + */ + public static UnionWithOperation unionWith(String collection) { + + Assert.notNull(collection, "Collection must not be null!"); + return new UnionWithOperation(collection, null, null); + } + + /** + * Set the {@link AggregationPipeline} to apply to the specified collection. The pipeline corresponds to the optional + * {@code pipeline} field of the {@code $unionWith} aggregation stage and is used to compute the documents going into + * the result set. + * + * @param pipeline the {@link AggregationPipeline} that computes the documents. Must not be {@literal null}. + * @return new instance of {@link UnionWithOperation}. + * @throws IllegalArgumentException if the required argument is {@literal null}. + */ + public UnionWithOperation pipeline(AggregationPipeline pipeline) { + return new UnionWithOperation(collection, pipeline, domainType); + } + + /** + * Set the aggregation pipeline stages to apply to the specified collection. The pipeline corresponds to the optional + * {@code pipeline} field of the {@code $unionWith} aggregation stage and is used to compute the documents going into + * the result set. + * + * @param aggregationStages the aggregation pipeline stages that compute the documents. Must not be {@literal null}. + * @return new instance of {@link UnionWithOperation}. + * @throws IllegalArgumentException if the required argument is {@literal null}. + */ + public UnionWithOperation pipeline(List aggregationStages) { + return new UnionWithOperation(collection, new AggregationPipeline(aggregationStages), domainType); + } + + /** + * Set the aggregation pipeline stages to apply to the specified collection. The pipeline corresponds to the optional + * {@code pipeline} field of the {@code $unionWith} aggregation stage and is used to compute the documents going into + * the result set. + * + * @param aggregationStages the aggregation pipeline stages that compute the documents. Must not be {@literal null}. + * @return new instance of {@link UnionWithOperation}. + * @throws IllegalArgumentException if the required argument is {@literal null}. + */ + public UnionWithOperation pipeline(AggregationOperation... aggregationStages) { + return new UnionWithOperation(collection, new AggregationPipeline(Arrays.asList(aggregationStages)), domainType); + } + + /** + * Set domain type used for field name mapping of property references used by the {@link AggregationPipeline}. + * Remember to also use a {@link TypedAggregation} in the outer pipeline.
+ * If not set, field names used within {@link AggregationOperation pipeline operations} are taken as is. + * + * @param domainType the domain type to map field names used in pipeline operations to. Must not be {@literal null}. + * @return new instance of {@link UnionWithOperation}. + * @throws IllegalArgumentException if the required argument is {@literal null}. + */ + public UnionWithOperation mapFieldsTo(Class domainType) { + + Assert.notNull(domainType, "DomainType must not be null!"); + return new UnionWithOperation(collection, pipeline, domainType); + } + + /* + * (non-Javadoc) + * @see org.springframework.data.mongodb.core.aggregation.AggregationOperation#toDocument(org.springframework.data.mongodb.core.aggregation.AggregationOperationContext) + */ + @Override + public Document toDocument(AggregationOperationContext context) { + + Document $unionWith = new Document("coll", collection); + if (pipeline == null || pipeline.isEmpty()) { + return new Document(getOperator(), $unionWith); + } + + $unionWith.append("pipeline", pipeline.toDocuments(computeContext(context))); + return new Document(getOperator(), $unionWith); + } + + private AggregationOperationContext computeContext(AggregationOperationContext source) { + + if (domainType == null) { + return Aggregation.DEFAULT_CONTEXT; + } + + if (source instanceof TypeBasedAggregationOperationContext) { + return ((TypeBasedAggregationOperationContext) source).continueOnMissingFieldReference(domainType); + } + + if (source instanceof ExposedFieldsAggregationOperationContext) { + return computeContext(((ExposedFieldsAggregationOperationContext) source).getRootContext()); + } + + return source; + } + + /* + * (non-Javadoc) + * @see org.springframework.data.mongodb.core.aggregation.AggregationOperation#getOperator() + */ + @Override + public String getOperator() { + return "$unionWith"; + } +} diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/UnionWithOperationUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/UnionWithOperationUnitTests.java new file mode 100644 index 000000000..880106bd0 --- /dev/null +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/UnionWithOperationUnitTests.java @@ -0,0 +1,129 @@ +/* + * 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 static org.assertj.core.api.Assertions.*; + +import java.util.Arrays; +import java.util.List; + +import org.bson.Document; +import org.junit.jupiter.api.Test; +import org.springframework.data.mongodb.core.convert.MappingMongoConverter; +import org.springframework.data.mongodb.core.convert.NoOpDbRefResolver; +import org.springframework.data.mongodb.core.convert.QueryMapper; +import org.springframework.data.mongodb.core.mapping.Field; +import org.springframework.data.mongodb.core.mapping.MongoMappingContext; +import org.springframework.lang.Nullable; + +/** + * @author Christoph Strobl + */ +class UnionWithOperationUnitTests { + + @Test // DATAMONGO-2622 + void throwsErrorWhenNoCollectionPresent() { + assertThatExceptionOfType(IllegalArgumentException.class).isThrownBy(() -> UnionWithOperation.unionWith(null)); + } + + @Test // DATAMONGO-2622 + void rendersJustCollectionCorrectly() { + + assertThat(UnionWithOperation.unionWith("coll-1").toPipelineStages(contextFor(Warehouse.class))) + .containsExactly(new Document("$unionWith", new Document("coll", "coll-1"))); + } + + @Test // DATAMONGO-2622 + void rendersPipelineCorrectly() { + + assertThat(UnionWithOperation.unionWith("coll-1").mapFieldsTo(Warehouse.class) + .pipeline(Aggregation.project().and("location").as("region")).toPipelineStages(contextFor(Warehouse.class))) + .containsExactly(new Document("$unionWith", new Document("coll", "coll-1").append("pipeline", + Arrays.asList(new Document("$project", new Document("region", 1)))))); + } + + @Test // DATAMONGO-2622 + void rendersPipelineCorrectlyForDifferentDomainType() { + + assertThat(UnionWithOperation.unionWith("coll-1").pipeline(Aggregation.project().and("name").as("name")) + .mapFieldsTo(Supplier.class).toPipelineStages(contextFor(Warehouse.class))) + .containsExactly(new Document("$unionWith", new Document("coll", "coll-1").append("pipeline", + Arrays.asList(new Document("$project", new Document("name", "$supplier")))))); + } + + @Test // DATAMONGO-2622 + void rendersPipelineCorrectlyForUntypedContext() { + + assertThat(UnionWithOperation.unionWith("coll-1").pipeline(Aggregation.project("region")) + .toPipelineStages(contextFor(null))) + .containsExactly(new Document("$unionWith", new Document("coll", "coll-1").append("pipeline", + Arrays.asList(new Document("$project", new Document("region", 1)))))); + } + + @Test // DATAMONGO-2622 + void doesNotMapAgainstFieldsFromAPreviousStage() { + + TypedAggregation agg = TypedAggregation.newAggregation(Supplier.class, + Aggregation.project().and("name").as("supplier"), + UnionWithOperation.unionWith("coll-1").pipeline(Aggregation.project().and("name").as("name"))); + + List pipeline = agg.toPipeline(contextFor(Supplier.class)); + System.out.println("pipeline: " + pipeline); + assertThat(pipeline).containsExactly(new Document("$project", new Document("supplier", 1)), // + new Document("$unionWith", new Document("coll", "coll-1").append("pipeline", + Arrays.asList(new Document("$project", new Document("name", 1)))))); + } + + @Test // DATAMONGO-2622 + void mapAgainstUnionWithDomainTypeEvenWhenInsideTypedAggregation() { + + TypedAggregation agg = TypedAggregation.newAggregation(Supplier.class, + Aggregation.project().and("name").as("supplier"), UnionWithOperation.unionWith("coll-1") + .mapFieldsTo(Warehouse.class).pipeline(Aggregation.project().and("location").as("location"))); + + List pipeline = agg.toPipeline(contextFor(Supplier.class)); + assertThat(pipeline).containsExactly(new Document("$project", new Document("supplier", 1)), // + new Document("$unionWith", new Document("coll", "coll-1").append("pipeline", + Arrays.asList(new Document("$project", new Document("location", "$region")))))); + } + + private static AggregationOperationContext contextFor(@Nullable Class type) { + + if (type == null) { + return Aggregation.DEFAULT_CONTEXT; + } + + MappingMongoConverter mongoConverter = new MappingMongoConverter(NoOpDbRefResolver.INSTANCE, + new MongoMappingContext()); + mongoConverter.afterPropertiesSet(); + + return new TypeBasedAggregationOperationContext(type, mongoConverter.getMappingContext(), + new QueryMapper(mongoConverter)); + } + + static class Warehouse { + + String name; + @Field("region") String location; + String state; + } + + static class Supplier { + + @Field("supplier") String name; + String state; + } +}