From fb499f8c698ec9763d0c4346a98031deec9971b4 Mon Sep 17 00:00:00 2001
From: Christoph Strobl <cstrobl@pivotal.io>
Date: Tue, 14 Jan 2020 10:51:08 +0100
Subject: [PATCH] DATAMONGO-2450 - Apply array filters to bulk write.

Original pull request: #820.
---
 .../mongodb/core/DefaultBulkOperations.java   | 31 +++++++++++++++----
 .../core/DefaultBulkOperationsUnitTests.java  | 15 ++++++++-
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/DefaultBulkOperations.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/DefaultBulkOperations.java
index a1e0f2155..9c78f5fe3 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/DefaultBulkOperations.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/DefaultBulkOperations.java
@@ -15,6 +15,9 @@
  */
 package org.springframework.data.mongodb.core;
 
+import lombok.NonNull;
+import lombok.Value;
+
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -38,6 +41,8 @@ import org.springframework.data.mongodb.core.mapping.event.MongoMappingEvent;
 import org.springframework.data.mongodb.core.query.Collation;
 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.util.Pair;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
@@ -47,9 +52,6 @@ import com.mongodb.bulk.BulkWriteResult;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.model.*;
 
-import lombok.NonNull;
-import lombok.Value;
-
 /**
  * Default implementation for {@link BulkOperations}.
  *
@@ -349,9 +351,7 @@ class DefaultBulkOperations implements BulkOperations {
 		Assert.notNull(query, "Query must not be null!");
 		Assert.notNull(update, "Update must not be null!");
 
-		UpdateOptions options = new UpdateOptions();
-		options.upsert(upsert);
-		query.getCollation().map(Collation::toMongoCollation).ifPresent(options::collation);
+		UpdateOptions options = computeUpdateOptions(query, update, upsert);
 
 		if (multi) {
 			addModel(update, new UpdateManyModel<>(query.getQueryObject(), update.getUpdateObject(), options));
@@ -489,6 +489,25 @@ class DefaultBulkOperations implements BulkOperations {
 		throw new IllegalStateException("BulkMode was null!");
 	}
 
+	/**
+	 * @param filterQuery The {@link Query} to read a potential {@link Collation} from. Must not be {@literal null}.
+	 * @param update The {@link Update} to apply
+	 * @param upsert flag to indicate if document should be upserted.
+	 * @return new instance of {@link UpdateOptions}.
+	 */
+	private static UpdateOptions computeUpdateOptions(Query filterQuery, UpdateDefinition update, boolean upsert) {
+
+		UpdateOptions options = new UpdateOptions();
+		options.upsert(upsert);
+
+		if (update.hasArrayFilters()) {
+			options.arrayFilters(update.getArrayFilters().stream().map(ArrayFilter::asDocument).collect(Collectors.toList()));
+		}
+
+		filterQuery.getCollation().map(Collation::toMongoCollation).ifPresent(options::collation);
+		return options;
+	}
+
 	/**
 	 * {@link BulkOperationContext} holds information about
 	 * {@link org.springframework.data.mongodb.core.BulkOperations.BulkMode} the entity in use as well as references to
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/DefaultBulkOperationsUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/DefaultBulkOperationsUnitTests.java
index 9d1de31cf..4c9933995 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/DefaultBulkOperationsUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/DefaultBulkOperationsUnitTests.java
@@ -16,7 +16,7 @@
 package org.springframework.data.mongodb.core;
 
 import static org.assertj.core.api.Assertions.*;
-import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.*;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyList;
@@ -60,6 +60,7 @@ import org.springframework.data.mongodb.core.mapping.event.BeforeSaveCallback;
 import org.springframework.data.mongodb.core.mapping.event.BeforeSaveEvent;
 import org.springframework.data.mongodb.core.query.BasicQuery;
 import org.springframework.data.mongodb.core.query.Collation;
+import org.springframework.data.mongodb.core.query.Criteria;
 import org.springframework.data.mongodb.core.query.Update;
 
 import com.mongodb.MongoWriteException;
@@ -320,6 +321,18 @@ public class DefaultBulkOperationsUnitTests {
 		verify(collection).withWriteConcern(eq(WriteConcern.MAJORITY));
 	}
 
+	@Test // DATAMONGO-2450
+	public void appliesArrayFilterWhenPresent() {
+
+		ops.updateOne(new BasicQuery("{}"), new Update().filterArray(Criteria.where("element").gte(100))).execute();
+
+		verify(collection).bulkWrite(captor.capture(), any());
+
+		UpdateOneModel<Document> updateModel = (UpdateOneModel<Document>) captor.getValue().get(0);
+		assertThat(updateModel.getOptions().getArrayFilters().get(0))
+				.isEqualTo(new org.bson.Document("element", new Document("$gte", 100)));
+	}
+
 	class SomeDomainType {
 
 		@Id String id;