DATAMONGO-2331 - Add support for Update with an aggregation pipeline.

Now the update methods exposed by (Reactive)MongoOperations also accept an Aggregation Pipeline via AggregationUpdate.

The update can consist of the following stages:

* AggregationUpdate.set(...).toValue(...) -> $set : { ... }
* AggregationUpdate.unset(...) -> $unset : [ ... ]
* AggregationUpdate.replaceWith(...) -> $replaceWith : { ... }

AggregationUpdate update = Aggregation.newUpdate()
    .set("average").toValue(ArithmeticOperators.valueOf("tests").avg())
    .set("grade").toValue(ConditionalOperators.switchCases(
        when(valueOf("average").greaterThanEqualToValue(90)).then("A"),
        when(valueOf("average").greaterThanEqualToValue(80)).then("B"),
        when(valueOf("average").greaterThanEqualToValue(70)).then("C"),
        when(valueOf("average").greaterThanEqualToValue(60)).then("D"))
        .defaultTo("F")
    );

template.update(Student.class)
    .apply(update)
    .all();

Author: christophstrobl

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ExecutableUpdateOperation.java

@@ -19,6 +19,7 @@
 
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
+import org.springframework.data.mongodb.core.query.UpdateDefinition;
 import org.springframework.lang.Nullable;
 
 import com.mongodb.client.result.UpdateResult;
@@ -150,14 +151,27 @@ interface TerminatingUpdate<T> extends TerminatingFindAndModify<T>, FindAndModif
 	 */
 	interface UpdateWithUpdate<T> {
 
+		/**
+		 * Set the {@link UpdateDefinition} to be applied.
+		 *
+		 * @param update must not be {@literal null}.
+		 * @return new instance of {@link TerminatingUpdate}.
+		 * @throws IllegalArgumentException if update is {@literal null}.
+		 */
+		TerminatingUpdate<T> apply(UpdateDefinition update);
+
 		/**
 		 * Set the {@link Update} to be applied.
 		 *
 		 * @param update must not be {@literal null}.
 		 * @return new instance of {@link TerminatingUpdate}.
 		 * @throws IllegalArgumentException if update is {@literal null}.
+		 * @deprecated since 2.3 in favor of {@link #apply(UpdateDefinition)}.
 		 */
-		TerminatingUpdate<T> apply(Update update);
+		@Deprecated
+		default TerminatingUpdate<T> apply(Update update) {
+			return apply((UpdateDefinition) update);
+		}
 
 		/**
 		 * Specify {@code replacement} object.

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ExecutableUpdateOperationSupport.java

@@ -21,7 +21,7 @@
 import lombok.experimental.FieldDefaults;
 
 import org.springframework.data.mongodb.core.query.Query;
-import org.springframework.data.mongodb.core.query.Update;
+import org.springframework.data.mongodb.core.query.UpdateDefinition;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
@@ -67,7 +67,7 @@ static class ExecutableUpdateSupport<T>
 		@NonNull MongoTemplate template;
 		@NonNull Class domainType;
 		Query query;
-		@Nullable Update update;
+		@Nullable UpdateDefinition update;
 		@Nullable String collection;
 		@Nullable FindAndModifyOptions findAndModifyOptions;
 		@Nullable FindAndReplaceOptions findAndReplaceOptions;
@@ -76,10 +76,10 @@ static class ExecutableUpdateSupport<T>
 
 		/*
 		 * (non-Javadoc)
-		 * @see org.springframework.data.mongodb.core.ExecutableUpdateOperation.UpdateWithUpdate#apply(Update)
+		 * @see org.springframework.data.mongodb.core.ExecutableUpdateOperation.UpdateWithUpdate#apply(org.springframework.data.mongodb.core.query.UpdateDefinition)
 		 */
 		@Override
-		public TerminatingUpdate<T> apply(Update update) {
+		public TerminatingUpdate<T> apply(UpdateDefinition update) {
 
 			Assert.notNull(update, "Update must not be null!");
 

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoOperations.java

@@ -68,7 +68,9 @@
 import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext;
 import org.springframework.data.mongodb.core.aggregation.AggregationOptions;
 import org.springframework.data.mongodb.core.aggregation.AggregationResults;
+import org.springframework.data.mongodb.core.aggregation.AggregationUpdate;
 import org.springframework.data.mongodb.core.aggregation.Fields;
+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.DbRefResolver;
@@ -108,7 +110,6 @@
 import org.springframework.data.mongodb.core.query.Meta.CursorOption;
 import org.springframework.data.mongodb.core.query.NearQuery;
 import org.springframework.data.mongodb.core.query.Query;
-import org.springframework.data.mongodb.core.query.Update;
 import org.springframework.data.mongodb.core.query.UpdateDefinition;
 import org.springframework.data.mongodb.core.query.UpdateDefinition.ArrayFilter;
 import org.springframework.data.mongodb.core.validation.Validator;
@@ -1043,25 +1044,25 @@ public <T> GeoResults<T> geoNear(NearQuery near, Class<?> domainType, String col
 
 	@Nullable
 	@Override
-	public <T> T findAndModify(Query query, Update update, Class<T> entityClass) {
+	public <T> T findAndModify(Query query, UpdateDefinition update, Class<T> entityClass) {
 		return findAndModify(query, update, new FindAndModifyOptions(), entityClass, getCollectionName(entityClass));
 	}
 
 	@Nullable
 	@Override
-	public <T> T findAndModify(Query query, Update update, Class<T> entityClass, String collectionName) {
+	public <T> T findAndModify(Query query, UpdateDefinition update, Class<T> entityClass, String collectionName) {
 		return findAndModify(query, update, new FindAndModifyOptions(), entityClass, collectionName);
 	}
 
 	@Nullable
 	@Override
-	public <T> T findAndModify(Query query, Update update, FindAndModifyOptions options, Class<T> entityClass) {
+	public <T> T findAndModify(Query query, UpdateDefinition update, FindAndModifyOptions options, Class<T> entityClass) {
 		return findAndModify(query, update, options, entityClass, getCollectionName(entityClass));
 	}
 
 	@Nullable
 	@Override
-	public <T> T findAndModify(Query query, Update update, FindAndModifyOptions options, Class<T> entityClass,
+	public <T> T findAndModify(Query query, UpdateDefinition update, FindAndModifyOptions options, Class<T> entityClass,
 			String collectionName) {
 
 		Assert.notNull(query, "Query must not be null!");
@@ -1561,53 +1562,54 @@ public Object doInCollection(MongoCollection<Document> collection) throws MongoE
 	}
 
 	@Override
-	public UpdateResult upsert(Query query, Update update, Class<?> entityClass) {
+	public UpdateResult upsert(Query query, UpdateDefinition update, Class<?> entityClass) {
 		return doUpdate(getCollectionName(entityClass), query, update, entityClass, true, false);
 	}
 
 	@Override
-	public UpdateResult upsert(Query query, Update update, String collectionName) {
+	public UpdateResult upsert(Query query, UpdateDefinition update, String collectionName) {
 		return doUpdate(collectionName, query, update, null, true, false);
 	}
 
 	@Override
-	public UpdateResult upsert(Query query, Update update, Class<?> entityClass, String collectionName) {
+	public UpdateResult upsert(Query query, UpdateDefinition update, Class<?> entityClass, String collectionName) {
 
 		Assert.notNull(entityClass, "EntityClass must not be null!");
 
 		return doUpdate(collectionName, query, update, entityClass, true, false);
 	}
 
 	@Override
-	public UpdateResult updateFirst(Query query, Update update, Class<?> entityClass) {
+	public UpdateResult updateFirst(Query query, UpdateDefinition update, Class<?> entityClass) {
 		return doUpdate(getCollectionName(entityClass), query, update, entityClass, false, false);
 	}
 
 	@Override
-	public UpdateResult updateFirst(final Query query, final Update update, final String collectionName) {
+	public UpdateResult updateFirst(final Query query, final UpdateDefinition update, final String collectionName) {
 		return doUpdate(collectionName, query, update, null, false, false);
 	}
 
 	@Override
-	public UpdateResult updateFirst(Query query, Update update, Class<?> entityClass, String collectionName) {
+	public UpdateResult updateFirst(Query query, UpdateDefinition update, Class<?> entityClass, String collectionName) {
 
 		Assert.notNull(entityClass, "EntityClass must not be null!");
 
 		return doUpdate(collectionName, query, update, entityClass, false, false);
 	}
 
 	@Override
-	public UpdateResult updateMulti(Query query, Update update, Class<?> entityClass) {
+	public UpdateResult updateMulti(Query query, UpdateDefinition update, Class<?> entityClass) {
 		return doUpdate(getCollectionName(entityClass), query, update, entityClass, false, true);
 	}
 
 	@Override
-	public UpdateResult updateMulti(final Query query, final Update update, String collectionName) {
+	public UpdateResult updateMulti(final Query query, final UpdateDefinition update, String collectionName) {
 		return doUpdate(collectionName, query, update, null, false, true);
 	}
 
 	@Override
-	public UpdateResult updateMulti(final Query query, final Update update, Class<?> entityClass, String collectionName) {
+	public UpdateResult updateMulti(final Query query, final UpdateDefinition update, Class<?> entityClass,
+			String collectionName) {
 
 		Assert.notNull(entityClass, "EntityClass must not be null!");
 
@@ -1622,24 +1624,52 @@ protected UpdateResult doUpdate(final String collectionName, final Query query,
 		Assert.notNull(query, "Query must not be null!");
 		Assert.notNull(update, "Update must not be null!");
 
-		return execute(collectionName, collection -> {
+		MongoPersistentEntity<?> entity = entityClass == null ? null : getPersistentEntity(entityClass);
+		increaseVersionForUpdateIfNecessary(entity, update);
 
-			MongoPersistentEntity<?> entity = entityClass == null ? null : getPersistentEntity(entityClass);
+		UpdateOptions opts = new UpdateOptions();
+		opts.upsert(upsert);
 
-			increaseVersionForUpdateIfNecessary(entity, update);
+		if (update.hasArrayFilters()) {
+			opts.arrayFilters(update.getArrayFilters().stream().map(ArrayFilter::asDocument).collect(Collectors.toList()));
+		}
 
-			UpdateOptions opts = new UpdateOptions();
-			opts.upsert(upsert);
+		Document queryObj = new Document();
 
-			if (update.hasArrayFilters()) {
-				opts.arrayFilters(update.getArrayFilters().stream().map(ArrayFilter::asDocument).collect(Collectors.toList()));
-			}
+		if (query != null) {
+			queryObj.putAll(queryMapper.getMappedObject(query.getQueryObject(), entity));
+		}
 
-			Document queryObj = new Document();
+		if (multi && update.isIsolated() && !queryObj.containsKey("$isolated")) {
+			queryObj.put("$isolated", 1);
+		}
 
-			if (query != null) {
-				queryObj.putAll(queryMapper.getMappedObject(query.getQueryObject(), entity));
-			}
+		if (update instanceof AggregationUpdate) {
+
+			AggregationOperationContext context = entityClass != null
+					? new RelaxedTypeBasedAggregationOperationContext(entityClass, mappingContext, queryMapper)
+					: Aggregation.DEFAULT_CONTEXT;
+
+			AggregationUpdate aUppdate = ((AggregationUpdate) update);
+			List<Document> pipeline = new AggregationUtil(queryMapper, mappingContext).createPipeline(aUppdate, context);
+
+			return execute(collectionName, collection -> {
+
+				MongoAction mongoAction = new MongoAction(writeConcern, MongoActionOperation.UPDATE, collectionName,
+						entityClass, update.getUpdateObject(), queryObj);
+				WriteConcern writeConcernToUse = prepareWriteConcern(mongoAction);
+
+				collection = writeConcernToUse != null ? collection.withWriteConcern(writeConcernToUse) : collection;
+
+				if (multi) {
+					return collection.updateMany(queryObj, pipeline, opts);
+				}
+
+				return collection.updateOne(queryObj, pipeline, opts);
+			});
+		}
+
+		return execute(collectionName, collection -> {
 
 			operations.forType(entityClass) //
 					.getCollation(query) //
@@ -1649,10 +1679,6 @@ protected UpdateResult doUpdate(final String collectionName, final Query query,
 			Document updateObj = update instanceof MappedUpdate ? update.getUpdateObject()
 					: updateMapper.getMappedObject(update.getUpdateObject(), entity);
 
-			if (multi && update.isIsolated() && !queryObj.containsKey("$isolated")) {
-				queryObj.put("$isolated", 1);
-			}
-
 			if (LOGGER.isDebugEnabled()) {
 				LOGGER.debug("Calling update using query: {} and update: {} in collection: {}", serializeToJsonSafely(queryObj),
 						serializeToJsonSafely(updateObj), collectionName);
@@ -2640,7 +2666,7 @@ protected <T> T doFindAndRemove(String collectionName, Document query, Document
 
 	@SuppressWarnings("ConstantConditions")
 	protected <T> T doFindAndModify(String collectionName, Document query, Document fields, Document sort,
-			Class<T> entityClass, Update update, @Nullable FindAndModifyOptions options) {
+			Class<T> entityClass, UpdateDefinition update, @Nullable FindAndModifyOptions options) {
 
 		EntityReader<? super T, Bson> readerToUse = this.mongoConverter;
 
@@ -2653,7 +2679,18 @@ protected <T> T doFindAndModify(String collectionName, Document query, Document
 		increaseVersionForUpdateIfNecessary(entity, update);
 
 		Document mappedQuery = queryMapper.getMappedObject(query, entity);
-		Document mappedUpdate = updateMapper.getMappedObject(update.getUpdateObject(), entity);
+
+		Object mappedUpdate = new Document();
+		if (update instanceof AggregationUpdate) {
+
+			AggregationOperationContext context = entityClass != null
+					? new RelaxedTypeBasedAggregationOperationContext(entityClass, mappingContext, queryMapper)
+					: Aggregation.DEFAULT_CONTEXT;
+
+			mappedUpdate = new AggregationUtil(queryMapper, mappingContext).createPipeline((Aggregation) update, context);
+		} else {
+			mappedUpdate = updateMapper.getMappedObject(update.getUpdateObject(), entity);
+		}
 
 		if (LOGGER.isDebugEnabled()) {
 			LOGGER.debug(
@@ -3027,11 +3064,11 @@ private static class FindAndModifyCallback implements CollectionCallback<Documen
 		private final Document query;
 		private final Document fields;
 		private final Document sort;
-		private final Document update;
+		private final Object update;
 		private final List<Document> arrayFilters;
 		private final FindAndModifyOptions options;
 
-		public FindAndModifyCallback(Document query, Document fields, Document sort, Document update,
+		public FindAndModifyCallback(Document query, Document fields, Document sort, Object update,
 				List<Document> arrayFilters, FindAndModifyOptions options) {
 			this.query = query;
 			this.fields = fields;
@@ -3059,7 +3096,12 @@ public Document doInCollection(MongoCollection<Document> collection) throws Mong
 				opts.arrayFilters(arrayFilters);
 			}
 
-			return collection.findOneAndUpdate(query, update, opts);
+			if (update instanceof Document) {
+				return collection.findOneAndUpdate(query, (Document) update, opts);
+			} else if (update instanceof List) {
+				return collection.findOneAndUpdate(query, (List<Document>) update, opts);
+			}
+			throw new IllegalArgumentException("doh - that does not work");
 		}
 	}