Forward maxTimeMsec from Query to FindOneAndUpdateOptions.

Closes spring-projects#4776

Author: sriram-1729

Diff for: spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java

@@ -181,6 +181,7 @@
  * @author Bartłomiej Mazur
  * @author Michael Krog
  * @author Jakub Zurawa
+ * @author Kinjarapu Sriram
  */
 public class MongoTemplate
 		implements MongoOperations, ApplicationContextAware, IndexOperationsProvider, ReadPreferenceAware {
@@ -1088,7 +1089,7 @@ public <T> T findAndModify(Query query, UpdateDefinition update, FindAndModifyOp
 			operations.forType(entityClass).getCollation(query).ifPresent(optionsToUse::collation);
 		}
 
-		return doFindAndModify(createDelegate(query), collectionName, query.getQueryObject(), query.getFieldsObject(),
+		return doFindAndModify(createDelegate(query), collectionName, query, query.getFieldsObject(),
 				getMappedSortObject(query, entityClass), entityClass, update, optionsToUse);
 	}
 
@@ -2718,7 +2719,7 @@ protected <T> T doFindAndRemove(CollectionPreparer collectionPreparer, String co
 	}
 
 	@SuppressWarnings("ConstantConditions")
-	protected <T> T doFindAndModify(CollectionPreparer collectionPreparer, String collectionName, Document query,
+	protected <T> T doFindAndModify(CollectionPreparer collectionPreparer, String collectionName, Query query,
 			Document fields, Document sort, Class<T> entityClass, UpdateDefinition update,
 			@Nullable FindAndModifyOptions options) {
 
@@ -2728,13 +2729,35 @@ protected <T> T doFindAndModify(CollectionPreparer collectionPreparer, String co
 
 		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(entityClass);
 
-		UpdateContext updateContext = queryOperations.updateSingleContext(update, query, false);
+		UpdateContext updateContext = queryOperations.updateSingleContext(update, query.getQueryObject(), false);
 		updateContext.increaseVersionForUpdateIfNecessary(entity);
 
 		Document mappedQuery = updateContext.getMappedQuery(entity);
 		Object mappedUpdate = updateContext.isAggregationUpdate() ? updateContext.getUpdatePipeline(entityClass)
 				: updateContext.getMappedUpdate(entity);
 
+		FindOneAndUpdateOptions opts = new FindOneAndUpdateOptions();
+		opts.sort(sort);
+		if (options.isUpsert()) {
+			opts.upsert(true);
+		}
+		opts.projection(fields);
+		if (options.isReturnNew()) {
+			opts.returnDocument(ReturnDocument.AFTER);
+		}
+		options.getCollation().map(Collation::toMongoCollation).ifPresent(opts::collation);
+		List<Document> arrayFilters = update.getArrayFilters().stream().map(ArrayFilter::asDocument).collect(Collectors.toList());
+		if (!arrayFilters.isEmpty()) {
+			opts.arrayFilters(arrayFilters);
+		}
+		Meta meta = query.getMeta();
+		if (meta.hasValues()) {
+
+			if (meta.hasMaxTime()) {
+				opts.maxTime(meta.getRequiredMaxTimeMsec(), TimeUnit.MILLISECONDS);
+			}
+		}
+
 		if (LOGGER.isDebugEnabled()) {
 			LOGGER.debug(String.format(
 					"findAndModify using query: %s fields: %s sort: %s for class: %s and update: %s in collection: %s",
@@ -2744,8 +2767,7 @@ protected <T> T doFindAndModify(CollectionPreparer collectionPreparer, String co
 		}
 
 		return executeFindOneInternal(
-				new FindAndModifyCallback(collectionPreparer, mappedQuery, fields, sort, mappedUpdate,
-						update.getArrayFilters().stream().map(ArrayFilter::asDocument).collect(Collectors.toList()), options),
+				new FindAndModifyCallback(collectionPreparer, mappedQuery, mappedUpdate, opts),
 				new ReadDocumentCallback<>(this.mongoConverter, entityClass, collectionName), collectionName);
 	}
 
@@ -3157,47 +3179,25 @@ private static class FindAndModifyCallback implements CollectionCallback<Documen
 
 		private final CollectionPreparer<MongoCollection<Document>> collectionPreparer;
 		private final Document query;
-		private final Document fields;
-		private final Document sort;
 		private final Object update;
-		private final List<Document> arrayFilters;
-		private final FindAndModifyOptions options;
+		private final FindOneAndUpdateOptions options;
 
 		FindAndModifyCallback(CollectionPreparer<MongoCollection<Document>> collectionPreparer, Document query,
-				Document fields, Document sort, Object update, List<Document> arrayFilters, FindAndModifyOptions options) {
+				Object update, FindOneAndUpdateOptions options) {
 
 			this.collectionPreparer = collectionPreparer;
 			this.query = query;
-			this.fields = fields;
-			this.sort = sort;
 			this.update = update;
-			this.arrayFilters = arrayFilters;
 			this.options = options;
 		}
 
 		@Override
 		public Document doInCollection(MongoCollection<Document> collection) throws MongoException, DataAccessException {
 
-			FindOneAndUpdateOptions opts = new FindOneAndUpdateOptions();
-			opts.sort(sort);
-			if (options.isUpsert()) {
-				opts.upsert(true);
-			}
-			opts.projection(fields);
-			if (options.isReturnNew()) {
-				opts.returnDocument(ReturnDocument.AFTER);
-			}
-
-			options.getCollation().map(Collation::toMongoCollation).ifPresent(opts::collation);
-
-			if (!arrayFilters.isEmpty()) {
-				opts.arrayFilters(arrayFilters);
-			}
-
 			if (update instanceof Document document) {
-				return collectionPreparer.prepare(collection).findOneAndUpdate(query, document, opts);
+				return collectionPreparer.prepare(collection).findOneAndUpdate(query, document, options);
 			} else if (update instanceof List) {
-				return collectionPreparer.prepare(collection).findOneAndUpdate(query, (List<Document>) update, opts);
+				return collectionPreparer.prepare(collection).findOneAndUpdate(query, (List<Document>) update, options);
 			}
 
 			throw new IllegalArgumentException(String.format("Using %s is not supported in findOneAndUpdate", update));

Diff for: spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java

@@ -141,6 +141,7 @@
  * @author Roman Puchkovskiy
  * @author Yadhukrishna S Pai
  * @author Jakub Zurawa
+ * @author Kinjarapu Sriram
  */
 @MockitoSettings(strictness = Strictness.LENIENT)
 public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
@@ -2536,6 +2537,16 @@ public WriteConcern resolve(MongoAction action) {
 		verify(collection).withWriteConcern(eq(WriteConcern.UNACKNOWLEDGED));
 	}
 
+	@Test // GH-4776
+	void findAndModifyConsidersMaxTimeMs() {
+
+		template.findAndModify(new BasicQuery("{ 'spring' : 'data-mongodb' }").maxTimeMsec(5000), new Update(), Human.class);
+
+		ArgumentCaptor<FindOneAndUpdateOptions> options = ArgumentCaptor.forClass(FindOneAndUpdateOptions.class);
+		verify(collection).findOneAndUpdate(any(Bson.class), any(Bson.class), options.capture());
+		assertThat(options.getValue().getMaxTime(TimeUnit.MILLISECONDS)).isEqualTo(5000);
+	}
+
 	class AutogenerateableId {
 
 		@Id BigInteger id;