Skip to content

Commit

Permalink
Forward maxTimeMsec from Query to FindOneAndUpdateOptions.
Browse files Browse the repository at this point in the history
  • Loading branch information
sriram-1729 committed Sep 5, 2024
1 parent 7fbd496 commit ca29374
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 31 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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) {

Expand All @@ -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",
Expand All @@ -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);
}

Expand Down Expand Up @@ -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));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit ca29374

Please sign in to comment.