GODRIVER-3285 Allow update to supply sort option. (#1797)

Author: qingyang-hu

internal/integration/unified/bulkwrite_helpers.go

@@ -102,6 +102,8 @@ func createBulkWriteModel(rawModel bson.Raw) (mongo.WriteModel, error) {
 				if err != nil {
 					return nil, fmt.Errorf("error creating update: %w", err)
 				}
+			case "sort":
+				uom.SetSort(val.Document())
 			case "upsert":
 				uom.SetUpsert(val.Boolean())
 			default:
@@ -249,6 +251,8 @@ func createBulkWriteModel(rawModel bson.Raw) (mongo.WriteModel, error) {
 					return nil, fmt.Errorf("error creating hint: %w", err)
 				}
 				rom.SetHint(hint)
+			case "sort":
+				rom.SetSort(val.Document())
 			case "replacement":
 				replacement = val.Document()
 			case "upsert":

internal/integration/unified/collection_operation_execution.go

@@ -1293,6 +1293,8 @@ func executeReplaceOne(ctx context.Context, operation *operation) (*operationRes
 				return nil, fmt.Errorf("error creating hint: %w", err)
 			}
 			opts.SetHint(hint)
+		case "sort":
+			opts.SetSort(val.Document())
 		case "replacement":
 			replacement = val.Document()
 		case "upsert":

internal/integration/unified/crud_helpers.go

@@ -123,6 +123,8 @@ func createUpdateOneArguments(args bson.Raw) (*updateArguments, *options.UpdateO
 			}
 		case "upsert":
 			opts.SetUpsert(val.Boolean())
+		case "sort":
+			opts.SetSort(val.Document())
 		default:
 			return nil, nil, fmt.Errorf("unrecognized update option %q", key)
 		}

mongo/bulk_write.go

@@ -340,6 +340,7 @@ func (bw *bulkWrite) runUpdate(ctx context.Context, batch bulkWriteBatch) (opera
 				filter:    converted.Filter,
 				update:    converted.Replacement,
 				hint:      converted.Hint,
+				sort:      converted.Sort,
 				collation: converted.Collation,
 				upsert:    converted.Upsert,
 			}.marshal(bw.collection.bsonOpts, bw.collection.registry)
@@ -349,6 +350,7 @@ func (bw *bulkWrite) runUpdate(ctx context.Context, batch bulkWriteBatch) (opera
 				filter:         converted.Filter,
 				update:         converted.Update,
 				hint:           converted.Hint,
+				sort:           converted.Sort,
 				arrayFilters:   converted.ArrayFilters,
 				collation:      converted.Collation,
 				upsert:         converted.Upsert,
@@ -420,6 +422,7 @@ type updateDoc struct {
 	filter         interface{}
 	update         interface{}
 	hint           interface{}
+	sort           interface{}
 	arrayFilters   []interface{}
 	collation      *options.Collation
 	upsert         *bool
@@ -446,6 +449,16 @@ func (doc updateDoc) marshal(bsonOpts *options.BSONOptions, registry *bson.Regis
 	if doc.multi {
 		updateDoc = bsoncore.AppendBooleanElement(updateDoc, "multi", doc.multi)
 	}
+	if doc.sort != nil {
+		if isUnorderedMap(doc.sort) {
+			return nil, ErrMapForOrderedArgument{"sort"}
+		}
+		s, err := marshal(doc.sort, bsonOpts, registry)
+		if err != nil {
+			return nil, err
+		}
+		updateDoc = bsoncore.AppendDocumentElement(updateDoc, "sort", s)
+	}
 
 	if doc.arrayFilters != nil {
 		reg := registry

mongo/bulk_write_models.go

@@ -125,6 +125,7 @@ type ReplaceOneModel struct {
 	Filter      interface{}
 	Replacement interface{}
 	Hint        interface{}
+	Sort        interface{}
 }
 
 // NewReplaceOneModel creates a new ReplaceOneModel.
@@ -173,6 +174,14 @@ func (rom *ReplaceOneModel) SetUpsert(upsert bool) *ReplaceOneModel {
 	return rom
 }
 
+// SetSort specifies which document the operation replaces if the query matches multiple documents. The first document
+// matched by the sort order will be replaced. This option is only valid for MongoDB versions >= 8.0. The driver will
+// return an error if the sort parameter is a multi-key map. The default value is nil.
+func (rom *ReplaceOneModel) SetSort(sort interface{}) *ReplaceOneModel {
+	rom.Sort = sort
+	return rom
+}
+
 func (*ReplaceOneModel) writeModel() {}
 
 // UpdateOneModel is used to update at most one document in a BulkWrite operation.
@@ -183,6 +192,7 @@ type UpdateOneModel struct {
 	Update       interface{}
 	ArrayFilters []interface{}
 	Hint         interface{}
+	Sort         interface{}
 }
 
 // NewUpdateOneModel creates a new UpdateOneModel.
@@ -238,6 +248,14 @@ func (uom *UpdateOneModel) SetUpsert(upsert bool) *UpdateOneModel {
 	return uom
 }
 
+// SetSort specifies which document the operation updates if the query matches multiple documents. The first document
+// matched by the sort order will be updated. This option is only valid for MongoDB versions >= 8.0. The driver will
+// return an error if the sort parameter is a multi-key map. The default value is nil.
+func (uom *UpdateOneModel) SetSort(sort interface{}) *UpdateOneModel {
+	uom.Sort = sort
+	return uom
+}
+
 func (*UpdateOneModel) writeModel() {}
 
 // UpdateManyModel is used to update multiple documents in a BulkWrite operation.

mongo/collection.go

@@ -610,6 +610,7 @@ func (coll *Collection) updateOrReplace(
 	multi bool,
 	expectedRr returnResult,
 	checkDollarKey bool,
+	sort interface{},
 	args *options.UpdateManyOptions,
 ) (*UpdateResult, error) {
 
@@ -623,6 +624,7 @@ func (coll *Collection) updateOrReplace(
 		filter:         filter,
 		update:         update,
 		hint:           args.Hint,
+		sort:           sort,
 		arrayFilters:   args.ArrayFilters,
 		collation:      args.Collation,
 		upsert:         args.Upsert,
@@ -775,7 +777,7 @@ func (coll *Collection) UpdateOne(
 		Let:                      args.Let,
 	}
 
-	return coll.updateOrReplace(ctx, f, update, false, rrOne, true, updateOptions)
+	return coll.updateOrReplace(ctx, f, update, false, rrOne, true, args.Sort, updateOptions)
 }
 
 // UpdateMany executes an update command to update documents in the collection.
@@ -811,7 +813,7 @@ func (coll *Collection) UpdateMany(
 		return nil, fmt.Errorf("failed to construct options from builder: %w", err)
 	}
 
-	return coll.updateOrReplace(ctx, f, update, true, rrMany, true, args)
+	return coll.updateOrReplace(ctx, f, update, true, rrMany, true, nil, args)
 }
 
 // ReplaceOne executes an update command to replace at most one document in the collection.
@@ -865,7 +867,7 @@ func (coll *Collection) ReplaceOne(
 		Comment:                  args.Comment,
 	}
 
-	return coll.updateOrReplace(ctx, f, r, false, rrOne, false, updateOptions)
+	return coll.updateOrReplace(ctx, f, r, false, rrOne, false, args.Sort, updateOptions)
 }
 
 // Aggregate executes an aggregate command against the collection and returns a cursor over the resulting documents.

mongo/options/replaceoptions.go

@@ -17,6 +17,7 @@ type ReplaceOptions struct {
 	Hint                     interface{}
 	Upsert                   *bool
 	Let                      interface{}
+	Sort                     interface{}
 }
 
 // ReplaceOptionsBuilder contains options to configure replace operations. Each
@@ -122,3 +123,18 @@ func (ro *ReplaceOptionsBuilder) SetLet(l interface{}) *ReplaceOptionsBuilder {
 
 	return ro
 }
+
+// SetSort sets the value for the Sort field. Specifies a document specifying which document should
+// be replaced if the filter used by the operation matches multiple documents in the collection. If
+// set, the first document in the sorted order will be replaced. This option is only valid for MongoDB
+// versions >= 8.0. The driver will return an error if the sort parameter is a multi-key map. The
+// default value is nil.
+func (ro *ReplaceOptionsBuilder) SetSort(s interface{}) *ReplaceOptionsBuilder {
+	ro.Opts = append(ro.Opts, func(opts *ReplaceOptions) error {
+		opts.Sort = s
+
+		return nil
+	})
+
+	return ro
+}

mongo/options/updateoptions.go

@@ -18,6 +18,7 @@ type UpdateOneOptions struct {
 	Hint                     interface{}
 	Upsert                   *bool
 	Let                      interface{}
+	Sort                     interface{}
 }
 
 // UpdateOneOptionsBuilder contains options to configure UpdateOne operations.
@@ -137,6 +138,21 @@ func (uo *UpdateOneOptionsBuilder) SetLet(l interface{}) *UpdateOneOptionsBuilde
 	return uo
 }
 
+// SetSort sets the value for the Sort field. Specifies a document specifying which document should
+// be updated if the filter used by the operation matches multiple documents in the collection. If
+// set, the first document in the sorted order will be updated. This option is only valid for MongoDB
+// versions >= 8.0. The driver will return an error if the sort parameter is a multi-key map. The
+// default value is nil.
+func (uo *UpdateOneOptionsBuilder) SetSort(s interface{}) *UpdateOneOptionsBuilder {
+	uo.Opts = append(uo.Opts, func(opts *UpdateOneOptions) error {
+		opts.Sort = s
+
+		return nil
+	})
+
+	return uo
+}
+
 // UpdateManyOptions represents arguments that can be used to configure UpdateMany
 // operations.
 //