Allow sort option in client bulk write.

Author: qingyang-hu

internal/integration/unified/client_operation_execution.go

@@ -328,6 +328,7 @@ func createClientUpdateOneModel(value bson.Raw) (*mongo.ClientBulkWrite, error)
 		Collation    *options.Collation
 		Hint         *bson.RawValue
 		Upsert       *bool
+		Sort         *bson.RawValue
 	}
 	err := bson.Unmarshal(value, &v)
 	if err != nil {
@@ -340,12 +341,17 @@ func createClientUpdateOneModel(value bson.Raw) (*mongo.ClientBulkWrite, error)
 			return nil, err
 		}
 	}
+	var sort interface{}
+	if v.Sort != nil {
+		sort = v.Sort.Document()
+	}
 	model := &mongo.ClientUpdateOneModel{
 		Filter:    v.Filter,
 		Update:    v.Update,
 		Collation: v.Collation,
 		Hint:      hint,
 		Upsert:    v.Upsert,
+		Sort:      sort,
 	}
 	if len(v.ArrayFilters) > 0 {
 		model.ArrayFilters = v.ArrayFilters
@@ -405,6 +411,7 @@ func createClientReplaceOneModel(value bson.Raw) (*mongo.ClientBulkWrite, error)
 		Collation   *options.Collation
 		Hint        *bson.RawValue
 		Upsert      *bool
+		Sort        *bson.RawValue
 	}
 	err := bson.Unmarshal(value, &v)
 	if err != nil {
@@ -417,6 +424,10 @@ func createClientReplaceOneModel(value bson.Raw) (*mongo.ClientBulkWrite, error)
 			return nil, err
 		}
 	}
+	var sort interface{}
+	if v.Sort != nil {
+		sort = v.Sort.Document()
+	}
 	ns := strings.SplitN(v.Namespace, ".", 2)
 	return &mongo.ClientBulkWrite{
 		Database:   ns[0],
@@ -427,6 +438,7 @@ func createClientReplaceOneModel(value bson.Raw) (*mongo.ClientBulkWrite, error)
 			Collation:   v.Collation,
 			Hint:        hint,
 			Upsert:      v.Upsert,
+			Sort:        sort,
 		},
 	}, nil
 }

mongo/client_bulk_write.go

@@ -315,6 +315,7 @@ func (mb *modelBatches) appendBatches(fn functionSet, dst []byte, maxCount, tota
 				arrayFilters:   model.ArrayFilters,
 				collation:      model.Collation,
 				upsert:         model.Upsert,
+				sort:           model.Sort,
 				multi:          false,
 				checkDollarKey: true,
 			}).marshal(mb.client.bsonOpts, mb.client.registry)
@@ -342,6 +343,7 @@ func (mb *modelBatches) appendBatches(fn functionSet, dst []byte, maxCount, tota
 				arrayFilters:   nil,
 				collation:      model.Collation,
 				upsert:         model.Upsert,
+				sort:           model.Sort,
 				multi:          false,
 				checkDollarKey: false,
 			}).marshal(mb.client.bsonOpts, mb.client.registry)
@@ -603,6 +605,7 @@ type clientUpdateDoc struct {
 	hint           interface{}
 	arrayFilters   []interface{}
 	collation      *options.Collation
+	sort           interface{}
 	upsert         *bool
 	multi          bool
 	checkDollarKey bool
@@ -657,6 +660,17 @@ func (d *clientUpdateDoc) marshal(bsonOpts *options.BSONOptions, registry *bson.
 		doc = bsoncore.AppendValueElement(doc, "hint", hintVal)
 	}
 
+	if d.sort != nil {
+		if isUnorderedMap(d.sort) {
+			return nil, ErrMapForOrderedArgument{"sort"}
+		}
+		sortVal, err := marshalValue(d.sort, bsonOpts, registry)
+		if err != nil {
+			return nil, err
+		}
+		doc = bsoncore.AppendValueElement(doc, "sort", sortVal)
+	}
+
 	return bsoncore.AppendDocumentEnd(doc, uidx)
 }
 

mongo/client_bulk_write_models.go

@@ -52,6 +52,7 @@ type ClientUpdateOneModel struct {
 	Update       interface{}
 	ArrayFilters []interface{}
 	Hint         interface{}
+	Sort         interface{}
 }
 
 // NewClientUpdateOneModel creates a new ClientUpdateOneModel.
@@ -105,6 +106,14 @@ func (uom *ClientUpdateOneModel) SetUpsert(upsert bool) *ClientUpdateOneModel {
 	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 *ClientUpdateOneModel) SetSort(sort interface{}) *ClientUpdateOneModel {
+	uom.Sort = sort
+	return uom
+}
+
 // ClientUpdateManyModel is used to update multiple documents in a client-level BulkWrite operation.
 //
 // See corresponding setter methods for documentation.
@@ -176,6 +185,7 @@ type ClientReplaceOneModel struct {
 	Filter      interface{}
 	Replacement interface{}
 	Hint        interface{}
+	Sort        interface{}
 }
 
 // NewClientReplaceOneModel creates a new ClientReplaceOneModel.
@@ -222,6 +232,14 @@ func (rom *ClientReplaceOneModel) SetUpsert(upsert bool) *ClientReplaceOneModel
 	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 *ClientReplaceOneModel) SetSort(sort interface{}) *ClientReplaceOneModel {
+	rom.Sort = sort
+	return rom
+}
+
 // ClientDeleteOneModel is used to delete at most one document in a client-level BulkWriteOperation.
 //
 // See corresponding setter methods for documentation.