Skip to content

Commit

Permalink
Allow sort option in client bulk write.
Browse files Browse the repository at this point in the history
  • Loading branch information
qingyang-hu committed Jan 24, 2025
1 parent 5051f3b commit 24db535
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 0 deletions.
12 changes: 12 additions & 0 deletions internal/integration/unified/client_operation_execution.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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 {
Expand All @@ -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],
Expand All @@ -427,6 +438,7 @@ func createClientReplaceOneModel(value bson.Raw) (*mongo.ClientBulkWrite, error)
Collation: v.Collation,
Hint: hint,
Upsert: v.Upsert,
Sort: sort,
},
}, nil
}
Expand Down
14 changes: 14 additions & 0 deletions mongo/client_bulk_write.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -603,6 +605,7 @@ type clientUpdateDoc struct {
hint interface{}
arrayFilters []interface{}
collation *options.Collation
sort interface{}
upsert *bool
multi bool
checkDollarKey bool
Expand Down Expand Up @@ -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)
}

Expand Down
18 changes: 18 additions & 0 deletions mongo/client_bulk_write_models.go
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ type ClientUpdateOneModel struct {
Update interface{}
ArrayFilters []interface{}
Hint interface{}
Sort interface{}
}

// NewClientUpdateOneModel creates a new ClientUpdateOneModel.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -176,6 +185,7 @@ type ClientReplaceOneModel struct {
Filter interface{}
Replacement interface{}
Hint interface{}
Sort interface{}
}

// NewClientReplaceOneModel creates a new ClientReplaceOneModel.
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit 24db535

Please sign in to comment.