Allow updateOne and replaceOne to supply sort option (#1585)

Allow updateOne and replaceOne to supply sort option

Adding sort option for updateOne and replaceOne commands, so if it matches more than one candidate document, the first one matched by sort order will be updated.

JAVA-5722

Author: joykim1005

driver-core/src/main/com/mongodb/client/model/ReplaceOptions.java

@@ -37,6 +37,7 @@ public class ReplaceOptions {
     private String hintString;
     private BsonValue comment;
     private Bson variables;
+    private Bson sort;
 
     /**
      * Returns true if a new document should be inserted if there are no matches to the query filter.  The default is false.
@@ -221,6 +222,43 @@ public ReplaceOptions let(final Bson variables) {
         return this;
     }
 
+    /**
+     * Gets the sort criteria to apply to the operation.
+     *
+     * <p>
+     * The sort criteria determines which document the operation replaces if the query matches multiple documents.
+     * The first document matched by the sort criteria will be replaced.
+     * The default is null, which means no specific sort criteria is applied.
+     *
+     * @return a document describing the sort criteria, or null if no sort is specified.
+     * @mongodb.driver.manual reference/method/db.collection.replaceOne/ Sort
+     * @mongodb.server.release 8.0
+     * @since 5.3
+     * @see #sort(Bson)
+     */
+    @Nullable
+    public Bson getSort() {
+        return sort;
+    }
+
+    /**
+     * Sets the sort criteria to apply to the operation. A null value means no sort criteria is set.
+     *
+     * <p>
+     * The sort criteria determines which document the operation replaces if the query matches multiple documents.
+     * The first document matched by the specified sort criteria will be replaced.
+     *
+     * @param sort the sort criteria, which may be null.
+     * @return this
+     * @mongodb.driver.manual reference/method/db.collection.replaceOne/ Sort
+     * @mongodb.server.release 8.0
+     * @since 5.3
+     */
+    public ReplaceOptions sort(@Nullable final Bson sort) {
+        this.sort = sort;
+        return this;
+    }
+
     @Override
     public String toString() {
         return "ReplaceOptions{"
@@ -231,6 +269,7 @@ public String toString() {
                 + ", hintString=" + hintString
                 + ", comment=" + comment
                 + ", let=" + variables
+                + ", sort=" + sort
                 + '}';
     }
 }

driver-core/src/main/com/mongodb/client/model/UpdateOptions.java

@@ -40,6 +40,7 @@ public class UpdateOptions {
     private String hintString;
     private BsonValue comment;
     private Bson variables;
+    private Bson sort;
 
     /**
      * Returns true if a new document should be inserted if there are no matches to the query filter.  The default is false.
@@ -256,6 +257,43 @@ public UpdateOptions let(final Bson variables) {
         return this;
     }
 
+    /**
+     * Gets the sort criteria to apply to the operation.
+     *
+     * <p>
+     * The sort criteria determines which document the operation updates if the query matches multiple documents.
+     * The first document matched by the sort criteria will be updated.
+     * The default is null, which means no specific sort criteria is applied.
+     *
+     * @return a document describing the sort criteria, or null if no sort is specified.
+     * @mongodb.driver.manual reference/method/db.collection.updateOne/ Sort
+     * @mongodb.server.release 8.0
+     * @since 5.3
+     * @see #sort(Bson)
+     */
+    @Nullable
+    public Bson getSort() {
+        return sort;
+    }
+
+    /**
+     * Sets the sort criteria to apply to the operation. A null value means no sort criteria is set.
+     *
+     * <p>
+     * The sort criteria determines which document the operation updates if the query matches multiple documents.
+     * The first document matched by the specified sort criteria will be updated.
+     *
+     * @param sort the sort criteria, which may be null.
+     * @return this
+     * @mongodb.driver.manual reference/method/db.collection.updateOne/ Sort
+     * @mongodb.server.release 8.0
+     * @since 5.3
+     */
+    public UpdateOptions sort(@Nullable final Bson sort) {
+        this.sort = sort;
+        return this;
+    }
+
     @Override
     public String toString() {
         return "UpdateOptions{"
@@ -267,6 +305,7 @@ public String toString() {
                 + ", hintString=" + hintString
                 + ", comment=" + comment
                 + ", let=" + variables
+                + ", sort=" + sort
                 + '}';
     }
 }

driver-core/src/main/com/mongodb/internal/bulk/UpdateRequest.java

@@ -40,6 +40,7 @@ public final class UpdateRequest extends WriteRequest {
     private List<BsonDocument> arrayFilters;
     @Nullable private BsonDocument hint;
     @Nullable private String hintString;
+    @Nullable private BsonDocument sort;
 
     public UpdateRequest(final BsonDocument filter, @Nullable final BsonValue update, final Type updateType) {
         if (updateType != Type.UPDATE && updateType != Type.REPLACE) {
@@ -128,5 +129,15 @@ public UpdateRequest hintString(@Nullable final String hint) {
         this.hintString = hint;
         return this;
     }
+
+    @Nullable
+    public BsonDocument getSort() {
+        return sort;
+    }
+
+    public UpdateRequest sort(@Nullable final BsonDocument sort) {
+        this.sort = sort;
+        return this;
+    }
 }
 

driver-core/src/main/com/mongodb/internal/connection/SplittablePayload.java

@@ -270,6 +270,11 @@ public void encode(final BsonWriter writer, final WriteRequestWithIndex writeReq
                 } else if (update.getHintString() != null) {
                     writer.writeString("hint", update.getHintString());
                 }
+                if (update.getSort() != null) {
+                    writer.writeName("sort");
+                    getCodec(assertNotNull(update.getSort())).encode(writer, assertNotNull(update.getSort()),
+                            EncoderContext.builder().build());
+                }
                 writer.writeEndDocument();
             } else {
                 DeleteRequest deleteRequest = (DeleteRequest) writeRequestWithIndex.getWriteRequest();

driver-core/src/main/com/mongodb/internal/operation/Operations.java

@@ -464,7 +464,8 @@ MixedBulkWriteOperation bulkWrite(final List<? extends WriteModel<? extends TDoc
                         .upsert(replaceOneModel.getReplaceOptions().isUpsert())
                         .collation(replaceOneModel.getReplaceOptions().getCollation())
                         .hint(toBsonDocument(replaceOneModel.getReplaceOptions().getHint()))
-                        .hintString(replaceOneModel.getReplaceOptions().getHintString());
+                        .hintString(replaceOneModel.getReplaceOptions().getHintString())
+                        .sort(toBsonDocument(replaceOneModel.getReplaceOptions().getSort()));
             } else if (writeModel instanceof UpdateOneModel) {
                 UpdateOneModel<TDocument> updateOneModel = (UpdateOneModel<TDocument>) writeModel;
                 BsonValue update = updateOneModel.getUpdate() != null ? toBsonDocument(updateOneModel.getUpdate())
@@ -475,7 +476,8 @@ MixedBulkWriteOperation bulkWrite(final List<? extends WriteModel<? extends TDoc
                         .collation(updateOneModel.getOptions().getCollation())
                         .arrayFilters(toBsonDocumentList(updateOneModel.getOptions().getArrayFilters()))
                         .hint(toBsonDocument(updateOneModel.getOptions().getHint()))
-                        .hintString(updateOneModel.getOptions().getHintString());
+                        .hintString(updateOneModel.getOptions().getHintString())
+                        .sort(toBsonDocument(updateOneModel.getOptions().getSort()));
             } else if (writeModel instanceof UpdateManyModel) {
                 UpdateManyModel<TDocument> updateManyModel = (UpdateManyModel<TDocument>) writeModel;
                 BsonValue update = updateManyModel.getUpdate() != null ? toBsonDocument(updateManyModel.getUpdate())

driver-core/src/test/unit/com/mongodb/client/model/UpdateOptionsSpecification.groovy

@@ -78,4 +78,12 @@ class UpdateOptionsSpecification extends Specification {
         where:
         hint << [null, '_id_']
     }
+
+    def 'should set sort'() {
+        expect:
+        new UpdateOptions().sort(sort).getSort() == sort
+
+        where:
+        sort << [null, new BsonDocument('_id', new BsonInt32(1))]
+    }
 }

driver-core/src/test/unit/com/mongodb/internal/operation/UpdateRequestSpecification.groovy

@@ -125,4 +125,13 @@ class UpdateRequestSpecification extends Specification {
         where:
         arrayFilters << [null, [], [new BsonDocument('a.b', new BsonInt32(42))]]
     }
+
+    def 'should set sort property'() {
+        expect:
+        new UpdateRequest(new BsonDocument(), new BsonDocument(), type).sort(sort).getSort() == sort
+
+        where:
+        type << [WriteRequest.Type.UPDATE, WriteRequest.Type.REPLACE]
+        sort << [null, new BsonDocument('_id', new BsonInt32(1))]
+    }
 }

driver-sync/src/test/functional/com/mongodb/client/unified/UnifiedCrudHelper.java

@@ -1159,6 +1159,9 @@ private UpdateOptions getUpdateOptions(final BsonDocument arguments) {
                 case "collation":
                     options.collation(asCollation(cur.getValue().asDocument()));
                     break;
+                case "sort":
+                    options.sort(cur.getValue().asDocument());
+                    break;
                 default:
                     throw new UnsupportedOperationException("Unsupported argument: " + cur.getKey());
             }
@@ -1193,6 +1196,9 @@ private ReplaceOptions getReplaceOptions(final BsonDocument arguments) {
                 case "collation":
                     options.collation(asCollation(cur.getValue().asDocument()));
                     break;
+                case "sort":
+                    options.sort(cur.getValue().asDocument());
+                    break;
                 default:
                     throw new UnsupportedOperationException("Unsupported argument: " + cur.getKey());
             }

driver-sync/src/test/functional/com/mongodb/client/unified/UnifiedTestModifications.java

@@ -135,15 +135,6 @@ public static void doSkips(final TestDef def) {
                 .test("crud", "findOneAndUpdate-hint-unacknowledged", "Unacknowledged findOneAndUpdate with hint document on 4.4+ server")
                 .test("crud", "findOneAndDelete-hint-unacknowledged", "Unacknowledged findOneAndDelete with hint string on 4.4+ server")
                 .test("crud", "findOneAndDelete-hint-unacknowledged", "Unacknowledged findOneAndDelete with hint document on 4.4+ server");
-        def.skipJira("https://jira.mongodb.org/browse/JAVA-5622")
-                .test("crud", "updateOne-sort", "UpdateOne with sort option")
-                .test("crud", "updateOne-sort", "updateOne with sort option unsupported (server-side error)")
-                .test("crud", "replaceOne-sort", "ReplaceOne with sort option")
-                .test("crud", "replaceOne-sort", "replaceOne with sort option unsupported (server-side error)")
-                .test("crud", "BulkWrite updateOne-sort", "BulkWrite updateOne with sort option")
-                .test("crud", "BulkWrite updateOne-sort", "BulkWrite updateOne with sort option unsupported (server-side error)")
-                .test("crud", "BulkWrite replaceOne-sort", "BulkWrite replaceOne with sort option")
-                .test("crud", "BulkWrite replaceOne-sort", "BulkWrite replaceOne with sort option unsupported (server-side error)");
 
         // gridfs
 

driver-sync/src/test/unit/com/mongodb/client/internal/MongoCollectionSpecification.groovy

@@ -813,15 +813,15 @@ class MongoCollectionSpecification extends Specification {
         def expectedOperation = { boolean upsert, WriteConcern wc, Boolean bypassValidation, Collation collation ->
             new MixedBulkWriteOperation(namespace,
                     [new UpdateRequest(new BsonDocument('a', new BsonInt32(1)), new BsonDocument('a', new BsonInt32(10)), REPLACE)
-                             .collation(collation).upsert(upsert).hint(hint).hintString(hintString)], true, wc, retryWrites)
+                             .collation(collation).upsert(upsert).hint(hint).hintString(hintString).sort(sort)], true, wc, retryWrites)
                     .bypassDocumentValidation(bypassValidation)
         }
         def replaceOneMethod = collection.&replaceOne
 
         when:
         def result = execute(replaceOneMethod, session, new Document('a', 1), new Document('a', 10),
                 new ReplaceOptions().upsert(true).bypassDocumentValidation(bypassDocumentValidation).collation(collation)
-                        .hint(hint).hintString(hintString))
+                        .hint(hint).hintString(hintString).sort(sort))
         executor.getClientSession() == session
         def operation = executor.getWriteOperation() as MixedBulkWriteOperation
 
@@ -830,15 +830,16 @@ class MongoCollectionSpecification extends Specification {
         result == expectedResult
 
         where:
-        [bypassDocumentValidation, modifiedCount, upsertedId, writeConcern, session, retryWrites, hint, hintString] << [
+        [bypassDocumentValidation, modifiedCount, upsertedId, writeConcern, session, retryWrites, hint, hintString, sort] << [
                 [null, true, false],
                 [1],
                 [null, new BsonInt32(42)],
                 [ACKNOWLEDGED, UNACKNOWLEDGED],
                 [null, Stub(ClientSession)],
                 [true, false],
                 [null, new BsonDocument('_id', new BsonInt32(1))],
-                [null, '_id_']
+                [null, '_id_'],
+                [null, new BsonDocument('_id', new BsonInt32(1))]
         ].combinations()
     }
 
@@ -880,11 +881,11 @@ class MongoCollectionSpecification extends Specification {
         def collection = new MongoCollectionImpl(namespace, Document, codecRegistry, readPreference, writeConcern,
                 retryWrites, true, readConcern, JAVA_LEGACY, null, TIMEOUT_SETTINGS, executor)
         def expectedOperation = { boolean upsert, WriteConcern wc, Boolean bypassDocumentValidation, Collation collation,
-                                  List<Bson> filters, BsonDocument hintDoc, String hintStr ->
+                                  List<Bson> filters, BsonDocument hintDoc, String hintStr, BsonDocument sortDoc ->
             new MixedBulkWriteOperation(namespace,
                     [new UpdateRequest(new BsonDocument('a', new BsonInt32(1)), new BsonDocument('a', new BsonInt32(10)), UPDATE)
                              .multi(false).upsert(upsert).collation(collation).arrayFilters(filters)
-                             .hint(hintDoc).hintString(hintStr)], true, wc, retryWrites)
+                             .hint(hintDoc).hintString(hintStr).sort(sortDoc)], true, wc, retryWrites)
                     .bypassDocumentValidation(bypassDocumentValidation)
         }
         def updateOneMethod = collection.&updateOne
@@ -894,29 +895,30 @@ class MongoCollectionSpecification extends Specification {
         def operation = executor.getWriteOperation() as MixedBulkWriteOperation
 
         then:
-        expect operation, isTheSameAs(expectedOperation(false, writeConcern, null, null, null, null, null))
+        expect operation, isTheSameAs(expectedOperation(false, writeConcern, null, null, null, null, null, null))
         executor.getClientSession() == session
         result == expectedResult
 
         when:
         result = execute(updateOneMethod, session, new Document('a', 1), new Document('a', 10),
                 new UpdateOptions().upsert(true).bypassDocumentValidation(true).collation(collation)
-                        .arrayFilters(arrayFilters).hint(hint).hintString(hintString))
+                        .arrayFilters(arrayFilters).hint(hint).hintString(hintString).sort(sort))
         operation = executor.getWriteOperation() as MixedBulkWriteOperation
 
         then:
-        expect operation, isTheSameAs(expectedOperation(true, writeConcern, true, collation, arrayFilters, hint, hintString))
+        expect operation, isTheSameAs(expectedOperation(true, writeConcern, true, collation, arrayFilters, hint, hintString, sort))
         executor.getClientSession() == session
         result == expectedResult
 
         where:
-        [writeConcern, arrayFilters, session, retryWrites, hint, hintString] << [
+        [writeConcern, arrayFilters, session, retryWrites, hint, hintString, sort] << [
                 [ACKNOWLEDGED, UNACKNOWLEDGED],
                 [null, [], [new BsonDocument('a.b', new BsonInt32(42))]],
                 [null, Stub(ClientSession)],
                 [true, false],
                 [null, new BsonDocument('_id', new BsonInt32(1))],
-                [null, '_id_']
+                [null, '_id_'],
+                [null, new BsonDocument('_id', new BsonInt32(1))]
         ].combinations()
     }