DATAMONGO-1986 - Polishing.
Refactor duplicated code into AggregationUtil. Original pull request: #564.
This commit is contained in:
@@ -0,0 +1,118 @@
|
||||
/*
|
||||
* Copyright 2018 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
|
||||
*
|
||||
* http://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;
|
||||
|
||||
import lombok.AllArgsConstructor;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.bson.Document;
|
||||
import org.springframework.data.mapping.context.MappingContext;
|
||||
import org.springframework.data.mongodb.core.aggregation.Aggregation;
|
||||
import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext;
|
||||
import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext;
|
||||
import org.springframework.data.mongodb.core.aggregation.TypedAggregation;
|
||||
import org.springframework.data.mongodb.core.convert.QueryMapper;
|
||||
import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
|
||||
import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
|
||||
import org.springframework.lang.Nullable;
|
||||
import org.springframework.util.ObjectUtils;
|
||||
|
||||
/**
|
||||
* Utility methods to map {@link org.springframework.data.mongodb.core.aggregation.Aggregation} pipeline definitions and
|
||||
* create type-bound {@link AggregationOperationContext}.
|
||||
*
|
||||
* @author Christoph Strobl
|
||||
* @author Mark Paluch
|
||||
* @since 2.0.8
|
||||
*/
|
||||
@AllArgsConstructor
|
||||
class AggregationUtil {
|
||||
|
||||
QueryMapper queryMapper;
|
||||
MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext;
|
||||
|
||||
/**
|
||||
* Prepare the {@link AggregationOperationContext} for a given aggregation by either returning the context itself it
|
||||
* is not {@literal null}, create a {@link TypeBasedAggregationOperationContext} if the aggregation contains type
|
||||
* information (is a {@link TypedAggregation}) or use the {@link Aggregation#DEFAULT_CONTEXT}.
|
||||
*
|
||||
* @param aggregation must not be {@literal null}.
|
||||
* @param context can be {@literal null}.
|
||||
* @return the root {@link AggregationOperationContext} to use.
|
||||
*/
|
||||
AggregationOperationContext prepareAggregationContext(Aggregation aggregation,
|
||||
@Nullable AggregationOperationContext context) {
|
||||
|
||||
if (context != null) {
|
||||
return context;
|
||||
}
|
||||
|
||||
if (aggregation instanceof TypedAggregation) {
|
||||
return new TypeBasedAggregationOperationContext(((TypedAggregation) aggregation).getInputType(), mappingContext,
|
||||
queryMapper);
|
||||
}
|
||||
|
||||
return Aggregation.DEFAULT_CONTEXT;
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract and map the aggregation pipeline into a {@link List} of {@link Document}.
|
||||
*
|
||||
* @param aggregation
|
||||
* @param context
|
||||
* @return
|
||||
*/
|
||||
Document createPipeline(String collectionName, Aggregation aggregation, AggregationOperationContext context) {
|
||||
|
||||
if (!ObjectUtils.nullSafeEquals(context, Aggregation.DEFAULT_CONTEXT)) {
|
||||
return aggregation.toDocument(collectionName, context);
|
||||
}
|
||||
|
||||
Document command = aggregation.toDocument(collectionName, context);
|
||||
command.put("pipeline", mapAggregationPipeline(command.get("pipeline", List.class)));
|
||||
|
||||
return command;
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract the command and map the aggregation pipeline.
|
||||
*
|
||||
* @param aggregation
|
||||
* @param context
|
||||
* @return
|
||||
*/
|
||||
Document createCommand(String collection, Aggregation aggregation, AggregationOperationContext context) {
|
||||
|
||||
Document command = aggregation.toDocument(collection, context);
|
||||
|
||||
if (!ObjectUtils.nullSafeEquals(context, Aggregation.DEFAULT_CONTEXT)) {
|
||||
return command;
|
||||
}
|
||||
|
||||
command.put("pipeline", mapAggregationPipeline(command.get("pipeline", List.class)));
|
||||
|
||||
return command;
|
||||
}
|
||||
|
||||
private List<Document> mapAggregationPipeline(List<Document> pipeline) {
|
||||
|
||||
return pipeline.stream().map(val -> queryMapper.getMappedObject(val, Optional.empty()))
|
||||
.collect(Collectors.toList());
|
||||
}
|
||||
}
|
||||
@@ -27,7 +27,6 @@ import java.io.IOException;
|
||||
import java.util.*;
|
||||
import java.util.Map.Entry;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.bson.Document;
|
||||
import org.bson.conversions.Bson;
|
||||
@@ -1934,7 +1933,8 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
|
||||
Assert.notNull(aggregation, "Aggregation pipeline must not be null!");
|
||||
Assert.notNull(outputType, "Output type must not be null!");
|
||||
|
||||
Document commandResult = new BatchAggregationLoader(this, readPreference, Integer.MAX_VALUE)
|
||||
Document commandResult = new BatchAggregationLoader(this, queryMapper, mappingContext, readPreference,
|
||||
Integer.MAX_VALUE)
|
||||
.aggregate(collectionName, aggregation, context);
|
||||
|
||||
return new AggregationResults<>(returnPotentiallyMappedResults(outputType, commandResult, collectionName),
|
||||
@@ -1957,9 +1957,9 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
|
||||
return Collections.emptyList();
|
||||
}
|
||||
|
||||
DocumentCallback<O> callback = new UnwrapAndReadDocumentCallback<O>(mongoConverter, outputType, collectionName);
|
||||
DocumentCallback<O> callback = new UnwrapAndReadDocumentCallback<>(mongoConverter, outputType, collectionName);
|
||||
|
||||
List<O> mappedResults = new ArrayList<O>();
|
||||
List<O> mappedResults = new ArrayList<>();
|
||||
for (Document document : resultSet) {
|
||||
mappedResults.add(callback.doWith(document));
|
||||
}
|
||||
@@ -1974,9 +1974,10 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
|
||||
Assert.notNull(aggregation, "Aggregation pipeline must not be null!");
|
||||
Assert.notNull(outputType, "Output type must not be null!");
|
||||
|
||||
AggregationOperationContext rootContext = context == null ? Aggregation.DEFAULT_CONTEXT : context;
|
||||
AggregationUtil aggregationUtil = new AggregationUtil(queryMapper, mappingContext);
|
||||
AggregationOperationContext rootContext = aggregationUtil.prepareAggregationContext(aggregation, context);
|
||||
|
||||
Document command = aggregation.toDocument(collectionName, rootContext);
|
||||
Document command = aggregationUtil.createCommand(collectionName, aggregation, rootContext);
|
||||
|
||||
assertNotExplain(command);
|
||||
|
||||
@@ -1984,7 +1985,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
|
||||
LOGGER.debug("Streaming aggregation: {}", serializeToJsonSafely(command));
|
||||
}
|
||||
|
||||
ReadDocumentCallback<O> readCallback = new ReadDocumentCallback<O>(mongoConverter, outputType, collectionName);
|
||||
ReadDocumentCallback<O> readCallback = new ReadDocumentCallback<>(mongoConverter, outputType, collectionName);
|
||||
|
||||
return execute(collectionName, new CollectionCallback<CloseableIterator<O>>() {
|
||||
|
||||
@@ -2008,7 +2009,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
|
||||
cursor = cursor.collation(options.getCollation().map(Collation::toMongoCollation).get());
|
||||
}
|
||||
|
||||
return new CloseableIterableCursorAdapter<O>(cursor.iterator(), exceptionTranslator, readCallback);
|
||||
return new CloseableIterableCursorAdapter<>(cursor.iterator(), exceptionTranslator, readCallback);
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -2577,72 +2578,6 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
|
||||
return fields;
|
||||
}
|
||||
|
||||
/**
|
||||
* Prepare the {@link AggregationOperationContext} for a given aggregation by either returning the context itself it
|
||||
* is not {@literal null}, create a {@link TypeBasedAggregationOperationContext} if the aggregation contains type
|
||||
* information (is a {@link TypedAggregation}) or use the {@link Aggregation#DEFAULT_CONTEXT}.
|
||||
*
|
||||
* @param aggregation must not be {@literal null}.
|
||||
* @param context can be {@literal null}.
|
||||
* @return the root {@link AggregationOperationContext} to use.
|
||||
*/
|
||||
private AggregationOperationContext prepareAggregationContext(Aggregation aggregation,
|
||||
@Nullable AggregationOperationContext context) {
|
||||
|
||||
if (context != null) {
|
||||
return context;
|
||||
}
|
||||
|
||||
if (aggregation instanceof TypedAggregation) {
|
||||
return new TypeBasedAggregationOperationContext(((TypedAggregation) aggregation).getInputType(), mappingContext,
|
||||
queryMapper);
|
||||
}
|
||||
|
||||
return Aggregation.DEFAULT_CONTEXT;
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract and map the aggregation pipeline.
|
||||
*
|
||||
* @param aggregation
|
||||
* @param context
|
||||
* @return
|
||||
*/
|
||||
private Document aggregationToPipeline(String inputCollectionName, Aggregation aggregation, AggregationOperationContext context) {
|
||||
|
||||
if (!ObjectUtils.nullSafeEquals(context, Aggregation.DEFAULT_CONTEXT)) {
|
||||
return aggregation.toDocument(inputCollectionName, context);
|
||||
}
|
||||
|
||||
return queryMapper.getMappedObject(aggregation.toDocument(inputCollectionName, context), Optional.empty());
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract the command and map the aggregation pipeline.
|
||||
*
|
||||
* @param aggregation
|
||||
* @param context
|
||||
* @return
|
||||
*/
|
||||
private Document aggregationToCommand(String collection, Aggregation aggregation,
|
||||
AggregationOperationContext context) {
|
||||
|
||||
Document command = aggregation.toDocument(collection, context);
|
||||
|
||||
if (!ObjectUtils.nullSafeEquals(context, Aggregation.DEFAULT_CONTEXT)) {
|
||||
return command;
|
||||
}
|
||||
|
||||
command.put("pipeline", mapAggregationPipeline(command.get("pipeline", List.class)));
|
||||
|
||||
return command;
|
||||
}
|
||||
|
||||
private List<Document> mapAggregationPipeline(List<Document> pipeline) {
|
||||
|
||||
return pipeline.stream().map(val -> queryMapper.getMappedObject(val, Optional.empty()))
|
||||
.collect(Collectors.toList());
|
||||
}
|
||||
|
||||
/**
|
||||
* Tries to convert the given {@link RuntimeException} into a {@link DataAccessException} but returns the original
|
||||
@@ -3157,12 +3092,18 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
|
||||
private static final String OK = "ok";
|
||||
|
||||
private final MongoTemplate template;
|
||||
private final QueryMapper queryMapper;
|
||||
private final MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext;
|
||||
private final ReadPreference readPreference;
|
||||
private final int batchSize;
|
||||
|
||||
BatchAggregationLoader(MongoTemplate template, ReadPreference readPreference, int batchSize) {
|
||||
BatchAggregationLoader(MongoTemplate template, QueryMapper queryMapper,
|
||||
MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext,
|
||||
ReadPreference readPreference, int batchSize) {
|
||||
|
||||
this.template = template;
|
||||
this.queryMapper = queryMapper;
|
||||
this.mappingContext = mappingContext;
|
||||
this.readPreference = readPreference;
|
||||
this.batchSize = batchSize;
|
||||
}
|
||||
@@ -3185,11 +3126,13 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
|
||||
* Pre process the aggregation command sent to the server by adding {@code cursor} options to match execution on
|
||||
* different server versions.
|
||||
*/
|
||||
private static Document prepareAggregationCommand(String collectionName, Aggregation aggregation,
|
||||
private Document prepareAggregationCommand(String collectionName, Aggregation aggregation,
|
||||
@Nullable AggregationOperationContext context, int batchSize) {
|
||||
|
||||
AggregationOperationContext rootContext = context == null ? Aggregation.DEFAULT_CONTEXT : context;
|
||||
Document command = aggregation.toDocument(collectionName, rootContext);
|
||||
AggregationUtil aggregationUtil = new AggregationUtil(queryMapper, mappingContext);
|
||||
|
||||
AggregationOperationContext rootContext = aggregationUtil.prepareAggregationContext(aggregation, context);
|
||||
Document command = aggregationUtil.createCommand(collectionName, aggregation, rootContext);
|
||||
|
||||
if (!aggregation.getOptions().isExplain()) {
|
||||
command.put(CURSOR_FIELD, new Document(BATCH_SIZE_FIELD, batchSize));
|
||||
|
||||
@@ -733,8 +733,9 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati
|
||||
Assert.hasText(collectionName, "Collection name must not be null or empty!");
|
||||
Assert.notNull(outputType, "Output type must not be null!");
|
||||
|
||||
AggregationOperationContext rootContext = prepareAggregationContext(aggregation, context);
|
||||
Document command = aggregationToPipeline(collectionName, aggregation, rootContext);
|
||||
AggregationUtil aggregationUtil = new AggregationUtil(queryMapper, mappingContext);
|
||||
AggregationOperationContext rootContext = aggregationUtil.prepareAggregationContext(aggregation, context);
|
||||
Document command = aggregationUtil.createPipeline(collectionName, aggregation, rootContext);
|
||||
AggregationOptions options = AggregationOptions.fromDocument(command);
|
||||
|
||||
Assert.isTrue(!options.isExplain(), "Cannot use explain option with streaming!");
|
||||
@@ -2197,46 +2198,6 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Prepare the {@link AggregationOperationContext} for a given aggregation by either returning the context itself it
|
||||
* is not {@literal null}, create a {@link TypeBasedAggregationOperationContext} if the aggregation contains type
|
||||
* information (is a {@link TypedAggregation}) or use the {@link Aggregation#DEFAULT_CONTEXT}.
|
||||
*
|
||||
* @param aggregation must not be {@literal null}.
|
||||
* @param context can be {@literal null}.
|
||||
* @return the root {@link AggregationOperationContext} to use.
|
||||
*/
|
||||
private AggregationOperationContext prepareAggregationContext(Aggregation aggregation,
|
||||
@Nullable AggregationOperationContext context) {
|
||||
|
||||
if (context != null) {
|
||||
return context;
|
||||
}
|
||||
|
||||
if (aggregation instanceof TypedAggregation) {
|
||||
return new TypeBasedAggregationOperationContext(((TypedAggregation) aggregation).getInputType(), mappingContext,
|
||||
queryMapper);
|
||||
}
|
||||
|
||||
return Aggregation.DEFAULT_CONTEXT;
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract and map the aggregation pipeline.
|
||||
*
|
||||
* @param aggregation
|
||||
* @param context
|
||||
* @return
|
||||
*/
|
||||
private Document aggregationToPipeline(String inputCollectionName, Aggregation aggregation, AggregationOperationContext context) {
|
||||
|
||||
if (!ObjectUtils.nullSafeEquals(context, Aggregation.DEFAULT_CONTEXT)) {
|
||||
return aggregation.toDocument(inputCollectionName, context);
|
||||
}
|
||||
|
||||
return queryMapper.getMappedObject(aggregation.toDocument(inputCollectionName, context), Optional.empty());
|
||||
}
|
||||
|
||||
/**
|
||||
* Tries to convert the given {@link RuntimeException} into a {@link DataAccessException} but returns the original
|
||||
* exception if the conversation failed. Thus allows safe re-throwing of the return value.
|
||||
|
||||
@@ -31,12 +31,16 @@ import org.mockito.junit.MockitoJUnitRunner;
|
||||
import org.springframework.data.mongodb.core.MongoTemplate.BatchAggregationLoader;
|
||||
import org.springframework.data.mongodb.core.aggregation.Aggregation;
|
||||
import org.springframework.data.mongodb.core.aggregation.TypedAggregation;
|
||||
import org.springframework.data.mongodb.core.convert.DbRefResolver;
|
||||
import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
|
||||
import org.springframework.data.mongodb.core.convert.QueryMapper;
|
||||
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
|
||||
|
||||
import com.mongodb.ReadPreference;
|
||||
|
||||
/**
|
||||
* Unit tests for {@link BatchAggregationLoader}.
|
||||
*
|
||||
*
|
||||
* @author Christoph Strobl
|
||||
* @author Mark Paluch
|
||||
*/
|
||||
@@ -47,6 +51,7 @@ public class BatchAggregationLoaderUnitTests {
|
||||
project().and("firstName").as("name"));
|
||||
|
||||
@Mock MongoTemplate template;
|
||||
@Mock DbRefResolver dbRefResolver;
|
||||
@Mock Document aggregationResult;
|
||||
@Mock Document getMoreResult;
|
||||
|
||||
@@ -60,7 +65,9 @@ public class BatchAggregationLoaderUnitTests {
|
||||
|
||||
@Before
|
||||
public void setUp() {
|
||||
loader = new BatchAggregationLoader(template, ReadPreference.primary(), 10);
|
||||
MongoMappingContext context = new MongoMappingContext();
|
||||
loader = new BatchAggregationLoader(template, new QueryMapper(new MappingMongoConverter(dbRefResolver, context)),
|
||||
context, ReadPreference.primary(), 10);
|
||||
}
|
||||
|
||||
@Test // DATAMONGO-1824
|
||||
@@ -84,7 +91,7 @@ public class BatchAggregationLoaderUnitTests {
|
||||
when(aggregationResult.get("cursor")).thenReturn(cursorWithoutMore);
|
||||
|
||||
Document result = loader.aggregate("person", AGGREGATION, Aggregation.DEFAULT_CONTEXT);
|
||||
|
||||
|
||||
assertThat((List) result.get("result")).contains(luke);
|
||||
|
||||
verify(template).executeCommand(any(Document.class), any(ReadPreference.class));
|
||||
|
||||
Reference in New Issue
Block a user