Skip to content

Commit

Permalink
GODRIVER-3285 Allow sort option in client bulk write. (#1923)
Browse files Browse the repository at this point in the history
  • Loading branch information
qingyang-hu authored Feb 11, 2025
1 parent a3ad820 commit eb76c5e
Show file tree
Hide file tree
Showing 12 changed files with 557 additions and 20 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
10 changes: 6 additions & 4 deletions mongo/bulk_write_models.go
Original file line number Diff line number Diff line change
Expand Up @@ -183,8 +183,9 @@ func (rom *ReplaceOneModel) SetUpsert(upsert bool) *ReplaceOneModel {
}

// 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.
// matched by the sort order will be replaced. This option is only valid for MongoDB versions >= 8.0. The sort parameter
// is evaluated sequentially, so the driver will return an error if it is a multi-key map (which is unordeded). The
// default value is nil.
func (rom *ReplaceOneModel) SetSort(sort interface{}) *ReplaceOneModel {
rom.Sort = sort
return rom
Expand Down Expand Up @@ -259,8 +260,9 @@ func (uom *UpdateOneModel) SetUpsert(upsert bool) *UpdateOneModel {
}

// 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.
// matched by the sort order will be updated. This option is only valid for MongoDB versions >= 8.0. The sort parameter
// is evaluated sequentially, so the driver will return an error if it is a multi-key map (which is unordeded). The
// default value is nil.
func (uom *UpdateOneModel) SetSort(sort interface{}) *UpdateOneModel {
uom.Sort = sort
return uom
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
20 changes: 20 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,15 @@ 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 sort parameter
// is evaluated sequentially, so the driver will return an error if it is a multi-key map (which is unordeded). 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 +186,7 @@ type ClientReplaceOneModel struct {
Filter interface{}
Replacement interface{}
Hint interface{}
Sort interface{}
}

// NewClientReplaceOneModel creates a new ClientReplaceOneModel.
Expand Down Expand Up @@ -222,6 +233,15 @@ 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 sort parameter
// is evaluated sequentially, so the driver will return an error if it is a multi-key map (which is unordeded). 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
24 changes: 14 additions & 10 deletions mongo/options/findoptions.go
Original file line number Diff line number Diff line change
Expand Up @@ -258,8 +258,8 @@ func (f *FindOptionsBuilder) SetSkip(i int64) *FindOptionsBuilder {
}

// SetSort sets the value for the Sort field. Sort is a document specifying the order in which
// documents should be returned. The driver will return an error if the sort parameter is a
// multi-key map.
// documents should be returned. The sort parameter is evaluated sequentially, so the driver will
// return an error if it is a multi-key map (which is unordeded). The default value is nil.
func (f *FindOptionsBuilder) SetSort(sort interface{}) *FindOptionsBuilder {
f.Opts = append(f.Opts, func(opts *FindOptions) error {
opts.Sort = sort
Expand Down Expand Up @@ -426,8 +426,9 @@ func (f *FindOneOptionsBuilder) SetSkip(i int64) *FindOneOptionsBuilder {
}

// SetSort sets the value for the Sort field. Sets a document specifying the sort order to
// apply to the query. The first document in the sorted order will be returned. The driver
// will return an error if the sort parameter is a multi-key map.
// apply to the query. The first document in the sorted order will be returned. The sort
// parameter is evaluated sequentially, so the driver will return an error if it is a multi-
// key map (which is unordeded). The default value is nil.
func (f *FindOneOptionsBuilder) SetSort(sort interface{}) *FindOneOptionsBuilder {
f.Opts = append(f.Opts, func(opts *FindOneOptions) error {
opts.Sort = sort
Expand Down Expand Up @@ -539,8 +540,9 @@ func (f *FindOneAndReplaceOptionsBuilder) SetReturnDocument(rd ReturnDocument) *

// SetSort sets the value for the Sort field. Sets 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. The driver will return an
// error if the sort parameter is a multi-key map. The default value is nil.
// If set, the first document in the sorted order will be replaced. The sort parameter is evaluated
// sequentially, so the driver will return an error if it is a multi-key map (which is unordeded).
// The default value is nil.
func (f *FindOneAndReplaceOptionsBuilder) SetSort(sort interface{}) *FindOneAndReplaceOptionsBuilder {
f.Opts = append(f.Opts, func(opts *FindOneAndReplaceOptions) error {
opts.Sort = sort
Expand Down Expand Up @@ -716,8 +718,9 @@ func (f *FindOneAndUpdateOptionsBuilder) SetReturnDocument(rd ReturnDocument) *F

// SetSort sets the value for the Sort field. Sets 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. The driver will return an
// error if the sort parameter is a multi-key map. The default value is nil.
// If set, the first document in the sorted order will be updated. The sort parameter is evaluated
// sequentially, so the driver will return an error if it is a multi-key map (which is unordeded).
// The default value is nil.
func (f *FindOneAndUpdateOptionsBuilder) SetSort(sort interface{}) *FindOneAndUpdateOptionsBuilder {
f.Opts = append(f.Opts, func(opts *FindOneAndUpdateOptions) error {
opts.Sort = sort
Expand Down Expand Up @@ -846,8 +849,9 @@ func (f *FindOneAndDeleteOptionsBuilder) SetProjection(projection interface{}) *

// SetSort sets the value for the Sort field. Sets 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 selected for replacement. The driver
// will return an error if the sort parameter is a multi-key map. The default value is nil.
// If set, the first document in the sorted order will be deleted. The sort parameter is evaluated
// sequentially, so the driver will return an error if it is a multi-key map (which is unordeded).
// The default value is nil.
func (f *FindOneAndDeleteOptionsBuilder) SetSort(sort interface{}) *FindOneAndDeleteOptionsBuilder {
f.Opts = append(f.Opts, func(opts *FindOneAndDeleteOptions) error {
opts.Sort = sort
Expand Down
5 changes: 3 additions & 2 deletions mongo/options/gridfsoptions.go
Original file line number Diff line number Diff line change
Expand Up @@ -328,8 +328,9 @@ func (f *GridFSFindOptionsBuilder) SetSkip(i int32) *GridFSFindOptionsBuilder {
}

// SetSort sets the value for the Sort field. Sets a document specifying the order
// in which documents should be returned. The driver will return an error if the
// sort parameter is a multi-key map.
// in which documents should be returned. The sort parameter is evaluated sequentially,
// so the driver will return an error if it is a multi-key map (which is unordeded).
// The default value is nil.
func (f *GridFSFindOptionsBuilder) SetSort(sort interface{}) *GridFSFindOptionsBuilder {
f.Opts = append(f.Opts, func(opts *GridFSFindOptions) error {
opts.Sort = sort
Expand Down
4 changes: 2 additions & 2 deletions mongo/options/replaceoptions.go
Original file line number Diff line number Diff line change
Expand Up @@ -127,8 +127,8 @@ func (ro *ReplaceOptionsBuilder) SetLet(l interface{}) *ReplaceOptionsBuilder {
// 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.
// versions >= 8.0. The sort parameter is evaluated sequentially, so the driver will return an error
// if it is a multi-key map (which is unordeded). The default value is nil.
func (ro *ReplaceOptionsBuilder) SetSort(s interface{}) *ReplaceOptionsBuilder {
ro.Opts = append(ro.Opts, func(opts *ReplaceOptions) error {
opts.Sort = s
Expand Down
4 changes: 2 additions & 2 deletions mongo/options/updateoptions.go
Original file line number Diff line number Diff line change
Expand Up @@ -141,8 +141,8 @@ func (uo *UpdateOneOptionsBuilder) SetLet(l interface{}) *UpdateOneOptionsBuilde
// 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.
// versions >= 8.0. The sort parameter is evaluated sequentially, so the driver will return an error
// if it is a multi-key map (which is unordeded). The default value is nil.
func (uo *UpdateOneOptionsBuilder) SetSort(s interface{}) *UpdateOneOptionsBuilder {
uo.Opts = append(uo.Opts, func(opts *UpdateOneOptions) error {
opts.Sort = s
Expand Down
Loading

0 comments on commit eb76c5e

Please sign in to comment.