From 178172be659654eb3f2e0a9150bd4348bf4d6c21 Mon Sep 17 00:00:00 2001 From: Felipe Gasper Date: Fri, 3 Jul 2026 09:23:40 -0400 Subject: [PATCH 1/3] AI-generated addCommandFields extensions --- mongo/client_bulk_write.go | 2 +- mongo/collection.go | 43 ++++- mongo/index_view.go | 10 ++ mongo/insert.go | 2 +- mongo/options/dropcollectionoptions.go | 6 + x/mongo/driver/operation/aggregate.go | 20 +++ x/mongo/driver/operation/count.go | 19 +++ x/mongo/driver/operation/create_indexes.go | 19 +++ x/mongo/driver/operation/delete.go | 19 +++ x/mongo/driver/operation/distinct.go | 20 +++ x/mongo/driver/operation/drop_collection.go | 21 +++ x/mongo/driver/operation/drop_indexes.go | 19 +++ x/mongo/driver/operation/find.go | 20 +++ x/mongo/driver/operation/find_and_modify.go | 2 +- x/mongo/driver/operation/list_indexes.go | 20 +++ x/mongo/driver/operation/update.go | 2 +- x/mongo/driver/xoptions/options.go | 111 +++++++++++++ x/mongo/driver/xoptions/options_test.go | 164 ++++++++++++++++++++ 18 files changed, 508 insertions(+), 11 deletions(-) diff --git a/mongo/client_bulk_write.go b/mongo/client_bulk_write.go index cdb52650f7..d8a663de14 100644 --- a/mongo/client_bulk_write.go +++ b/mongo/client_bulk_write.go @@ -158,7 +158,7 @@ func (bw *clientBulkWrite) newCommand() func([]byte, description.SelectedServer) if len(bw.additionalCmd) > 0 { doc, err := bson.Marshal(bw.additionalCmd) if err != nil { - return nil, err + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) } dst = append(dst, doc[4:len(doc)-1]...) } diff --git a/mongo/collection.go b/mongo/collection.go index 11c649cec0..b85cdfa46f 100644 --- a/mongo/collection.go +++ b/mongo/collection.go @@ -591,6 +591,9 @@ func (coll *Collection) delete( if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { op = op.RawData(rawData) } + if additionalCmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { + op = op.AdditionalCmd(additionalCmd) + } rr, err := processWriteError(op.Execute(ctx)) if rr&expectedRr == 0 { @@ -1214,6 +1217,9 @@ func (coll *Collection) CountDocuments(ctx context.Context, filter any, if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { op = op.RawData(rawData) } + if additionalCmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { + op = op.AdditionalCmd(additionalCmd) + } err = op.Execute(ctx) if err != nil { @@ -1302,6 +1308,9 @@ func (coll *Collection) EstimatedDocumentCount( if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { op = op.RawData(rawData) } + if additionalCmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { + op = op.AdditionalCmd(additionalCmd) + } err = op.Execute(ctx) return op.Result().N, wrapErrors(err) @@ -1397,6 +1406,9 @@ func (coll *Collection) Distinct( if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { op = op.RawData(rawData) } + if additionalCmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { + op = op.AdditionalCmd(additionalCmd) + } err = op.Execute(ctx) if err != nil { @@ -1605,6 +1617,9 @@ func (coll *Collection) find( if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { op = op.RawData(rawData) } + if additionalCmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { + op = op.AdditionalCmd(additionalCmd) + } if err = op.Execute(ctx); err != nil { return nil, wrapErrors(err) @@ -1805,6 +1820,9 @@ func (coll *Collection) FindOneAndDelete( if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { op = op.RawData(rawData) } + if additionalCmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { + op = op.AdditionalCmd(additionalCmd) + } return coll.findAndModify(ctx, op) } @@ -2081,6 +2099,11 @@ func (coll *Collection) Drop(ctx context.Context, opts ...options.Lister[options ef := args.EncryptedFields + var additionalCmd bson.D + if cmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { + additionalCmd = cmd + } + if ef == nil { ef = coll.db.getEncryptedFieldsFromMap(coll.name) } @@ -2093,14 +2116,17 @@ func (coll *Collection) Drop(ctx context.Context, opts ...options.Lister[options } if ef != nil { - return coll.dropEncryptedCollection(ctx, ef) + return coll.dropEncryptedCollection(ctx, ef, additionalCmd) } - return coll.drop(ctx) + return coll.drop(ctx, additionalCmd) } // dropEncryptedCollection drops a collection with EncryptedFields. -func (coll *Collection) dropEncryptedCollection(ctx context.Context, ef any) error { +// +// additionalCmd is applied only to the drop of the data collection, not to the +// associated encryption state collections. +func (coll *Collection) dropEncryptedCollection(ctx context.Context, ef any, additionalCmd bson.D) error { efBSON, err := marshal(ef, coll.bsonOpts, coll.registry) if err != nil { return fmt.Errorf("error transforming document: %w", err) @@ -2112,7 +2138,7 @@ func (coll *Collection) dropEncryptedCollection(ctx context.Context, ef any) err if err != nil { return err } - if err := coll.db.Collection(escCollection).drop(ctx); err != nil { + if err := coll.db.Collection(escCollection).drop(ctx, nil); err != nil { return err } @@ -2121,16 +2147,16 @@ func (coll *Collection) dropEncryptedCollection(ctx context.Context, ef any) err if err != nil { return err } - if err := coll.db.Collection(ecocCollection).drop(ctx); err != nil { + if err := coll.db.Collection(ecocCollection).drop(ctx, nil); err != nil { return err } // Drop the data collection. - return coll.drop(ctx) + return coll.drop(ctx, additionalCmd) } // drop drops a collection without EncryptedFields. -func (coll *Collection) drop(ctx context.Context) error { +func (coll *Collection) drop(ctx context.Context, additionalCmd bson.D) error { if ctx == nil { ctx = context.Background() } @@ -2163,6 +2189,9 @@ func (coll *Collection) drop(ctx context.Context) error { Deployment(coll.client.deployment).Crypt(coll.client.cryptFLE). ServerAPI(coll.client.serverAPI).Timeout(coll.client.timeout). Authenticator(coll.client.authenticator) + if len(additionalCmd) > 0 { + op = op.AdditionalCmd(additionalCmd) + } err = op.Execute(ctx) // ignore namespace not found errors diff --git a/mongo/index_view.go b/mongo/index_view.go index 921635851b..e35e50ff8b 100644 --- a/mongo/index_view.go +++ b/mongo/index_view.go @@ -13,6 +13,7 @@ import ( "fmt" "strconv" + "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/internal/mongoutil" "go.mongodb.org/mongo-driver/v2/internal/optionsutil" "go.mongodb.org/mongo-driver/v2/internal/serverselector" @@ -112,6 +113,9 @@ func (iv IndexView) List(ctx context.Context, opts ...options.Lister[options.Lis if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { op = op.RawData(rawData) } + if additionalCmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { + op = op.AdditionalCmd(additionalCmd) + } err = op.Execute(ctx) if err != nil { @@ -295,6 +299,9 @@ func (iv IndexView) CreateMany( if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { op = op.RawData(rawData) } + if additionalCmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { + op = op.AdditionalCmd(additionalCmd) + } _, err = processWriteError(op.Execute(ctx)) if err != nil { @@ -435,6 +442,9 @@ func (iv IndexView) drop(ctx context.Context, index any, opts ...options.Lister[ if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { op = op.RawData(rawData) } + if additionalCmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { + op = op.AdditionalCmd(additionalCmd) + } err = op.Execute(ctx) if err != nil { diff --git a/mongo/insert.go b/mongo/insert.go index ba05b16113..5652de23f6 100644 --- a/mongo/insert.go +++ b/mongo/insert.go @@ -139,7 +139,7 @@ func (i *insert) command(dst []byte, desc description.SelectedServer) ([]byte, e if len(i.additionalCmd) > 0 { doc, err := bson.Marshal(i.additionalCmd) if err != nil { - return nil, err + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) } dst = append(dst, doc[4:len(doc)-1]...) } diff --git a/mongo/options/dropcollectionoptions.go b/mongo/options/dropcollectionoptions.go index 013b4bbb77..1e88007ac7 100644 --- a/mongo/options/dropcollectionoptions.go +++ b/mongo/options/dropcollectionoptions.go @@ -6,12 +6,18 @@ package options +import "go.mongodb.org/mongo-driver/v2/internal/optionsutil" + // DropCollectionOptions represents arguments that can be used to configure a // Drop operation. // // See corresponding setter methods for documentation. type DropCollectionOptions struct { EncryptedFields any + + // Deprecated: This option is for internal use only and should not be set. It may be changed or removed in any + // release. + Internal optionsutil.Options } // DropCollectionOptionsBuilder contains options to configure collection drop diff --git a/x/mongo/driver/operation/aggregate.go b/x/mongo/driver/operation/aggregate.go index 22dbb8087f..94baa8ed3d 100644 --- a/x/mongo/driver/operation/aggregate.go +++ b/x/mongo/driver/operation/aggregate.go @@ -9,8 +9,10 @@ package operation import ( "context" "errors" + "fmt" "time" + "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/event" "go.mongodb.org/mongo-driver/v2/internal/driverutil" "go.mongodb.org/mongo-driver/v2/mongo/readconcern" @@ -53,6 +55,7 @@ type Aggregate struct { timeout *time.Duration omitMaxTimeMS bool rawData *bool + additionalCmd bson.D result driver.CursorResponse } @@ -161,6 +164,13 @@ func (a *Aggregate) command(dst []byte, desc description.SelectedServer) ([]byte if a.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) { dst = bsoncore.AppendBooleanElement(dst, "rawData", *a.rawData) } + if len(a.additionalCmd) > 0 { + doc, err := bson.Marshal(a.additionalCmd) + if err != nil { + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) + } + dst = append(dst, doc[4:len(doc)-1]...) + } for optionName, optionValue := range a.customOptions { dst = bsoncore.AppendValueElement(dst, optionName, optionValue) } @@ -465,3 +475,13 @@ func (a *Aggregate) RawData(rawData bool) *Aggregate { a.rawData = &rawData return a } + +// AdditionalCmd sets additional command fields to be attached. +func (a *Aggregate) AdditionalCmd(d bson.D) *Aggregate { + if a == nil { + a = new(Aggregate) + } + + a.additionalCmd = d + return a +} diff --git a/x/mongo/driver/operation/count.go b/x/mongo/driver/operation/count.go index 69e585ebfe..277e21b710 100644 --- a/x/mongo/driver/operation/count.go +++ b/x/mongo/driver/operation/count.go @@ -12,6 +12,7 @@ import ( "fmt" "time" + "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/event" "go.mongodb.org/mongo-driver/v2/internal/driverutil" "go.mongodb.org/mongo-driver/v2/mongo/readconcern" @@ -44,6 +45,7 @@ type Count struct { serverAPI *driver.ServerAPIOptions timeout *time.Duration rawData *bool + additionalCmd bson.D } // CountResult represents a count result returned by the server. @@ -155,6 +157,13 @@ func (c *Count) command(dst []byte, desc description.SelectedServer) ([]byte, er if c.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) { dst = bsoncore.AppendBooleanElement(dst, "rawData", *c.rawData) } + if len(c.additionalCmd) > 0 { + doc, err := bson.Marshal(c.additionalCmd) + if err != nil { + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) + } + dst = append(dst, doc[4:len(doc)-1]...) + } return dst, nil } @@ -350,3 +359,13 @@ func (c *Count) RawData(rawData bool) *Count { c.rawData = &rawData return c } + +// AdditionalCmd sets additional command fields to be attached. +func (c *Count) AdditionalCmd(d bson.D) *Count { + if c == nil { + c = new(Count) + } + + c.additionalCmd = d + return c +} diff --git a/x/mongo/driver/operation/create_indexes.go b/x/mongo/driver/operation/create_indexes.go index e6a0e2c0ee..10989227b5 100644 --- a/x/mongo/driver/operation/create_indexes.go +++ b/x/mongo/driver/operation/create_indexes.go @@ -12,6 +12,7 @@ import ( "fmt" "time" + "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/event" "go.mongodb.org/mongo-driver/v2/internal/driverutil" "go.mongodb.org/mongo-driver/v2/mongo/writeconcern" @@ -41,6 +42,7 @@ type CreateIndexes struct { serverAPI *driver.ServerAPIOptions timeout *time.Duration rawData *bool + additionalCmd bson.D } // CreateIndexesResult represents a createIndexes result returned by the server. @@ -142,6 +144,13 @@ func (ci *CreateIndexes) command(dst []byte, desc description.SelectedServer) ([ if ci.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) { dst = bsoncore.AppendBooleanElement(dst, "rawData", *ci.rawData) } + if len(ci.additionalCmd) > 0 { + doc, err := bson.Marshal(ci.additionalCmd) + if err != nil { + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) + } + dst = append(dst, doc[4:len(doc)-1]...) + } return dst, nil } @@ -318,3 +327,13 @@ func (ci *CreateIndexes) RawData(rawData bool) *CreateIndexes { ci.rawData = &rawData return ci } + +// AdditionalCmd sets additional command fields to be attached. +func (ci *CreateIndexes) AdditionalCmd(d bson.D) *CreateIndexes { + if ci == nil { + ci = new(CreateIndexes) + } + + ci.additionalCmd = d + return ci +} diff --git a/x/mongo/driver/operation/delete.go b/x/mongo/driver/operation/delete.go index 6a5209146d..38ff46a0a8 100644 --- a/x/mongo/driver/operation/delete.go +++ b/x/mongo/driver/operation/delete.go @@ -12,6 +12,7 @@ import ( "fmt" "time" + "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/event" "go.mongodb.org/mongo-driver/v2/internal/driverutil" "go.mongodb.org/mongo-driver/v2/internal/logger" @@ -46,6 +47,7 @@ type Delete struct { let bsoncore.Document timeout *time.Duration rawData *bool + additionalCmd bson.D logger *logger.Logger } @@ -148,6 +150,13 @@ func (d *Delete) command(dst []byte, desc description.SelectedServer) ([]byte, e if d.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) { dst = bsoncore.AppendBooleanElement(dst, "rawData", *d.rawData) } + if len(d.additionalCmd) > 0 { + doc, err := bson.Marshal(d.additionalCmd) + if err != nil { + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) + } + dst = append(dst, doc[4:len(doc)-1]...) + } return dst, nil } @@ -378,3 +387,13 @@ func (d *Delete) RawData(rawData bool) *Delete { d.rawData = &rawData return d } + +// AdditionalCmd sets additional command fields to be attached. +func (d *Delete) AdditionalCmd(cmd bson.D) *Delete { + if d == nil { + d = new(Delete) + } + + d.additionalCmd = cmd + return d +} diff --git a/x/mongo/driver/operation/distinct.go b/x/mongo/driver/operation/distinct.go index 411eb3cac3..e134e01894 100644 --- a/x/mongo/driver/operation/distinct.go +++ b/x/mongo/driver/operation/distinct.go @@ -9,8 +9,10 @@ package operation import ( "context" "errors" + "fmt" "time" + "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/event" "go.mongodb.org/mongo-driver/v2/internal/driverutil" "go.mongodb.org/mongo-driver/v2/mongo/readconcern" @@ -46,6 +48,7 @@ type Distinct struct { serverAPI *driver.ServerAPIOptions timeout *time.Duration rawData *bool + additionalCmd bson.D } // DistinctResult represents a distinct result returned by the server. @@ -138,6 +141,13 @@ func (d *Distinct) command(dst []byte, desc description.SelectedServer) ([]byte, if d.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) { dst = bsoncore.AppendBooleanElement(dst, "rawData", *d.rawData) } + if len(d.additionalCmd) > 0 { + doc, err := bson.Marshal(d.additionalCmd) + if err != nil { + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) + } + dst = append(dst, doc[4:len(doc)-1]...) + } return dst, nil } @@ -363,3 +373,13 @@ func (d *Distinct) RawData(rawData bool) *Distinct { d.rawData = &rawData return d } + +// AdditionalCmd sets additional command fields to be attached. +func (d *Distinct) AdditionalCmd(cmd bson.D) *Distinct { + if d == nil { + d = new(Distinct) + } + + d.additionalCmd = cmd + return d +} diff --git a/x/mongo/driver/operation/drop_collection.go b/x/mongo/driver/operation/drop_collection.go index a3880d9103..2c8ef9509e 100644 --- a/x/mongo/driver/operation/drop_collection.go +++ b/x/mongo/driver/operation/drop_collection.go @@ -12,6 +12,7 @@ import ( "fmt" "time" + "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/event" "go.mongodb.org/mongo-driver/v2/internal/driverutil" "go.mongodb.org/mongo-driver/v2/mongo/writeconcern" @@ -36,6 +37,7 @@ type DropCollection struct { result DropCollectionResult serverAPI *driver.ServerAPIOptions timeout *time.Duration + additionalCmd bson.D } // DropCollectionResult represents a dropCollection result returned by the server. @@ -112,9 +114,28 @@ func (dc *DropCollection) Execute(ctx context.Context) error { func (dc *DropCollection) command(dst []byte, _ description.SelectedServer) ([]byte, error) { dst = bsoncore.AppendStringElement(dst, "drop", dc.collection) + + if len(dc.additionalCmd) > 0 { + doc, err := bson.Marshal(dc.additionalCmd) + if err != nil { + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) + } + dst = append(dst, doc[4:len(doc)-1]...) + } + return dst, nil } +// AdditionalCmd sets additional command fields to be attached. +func (dc *DropCollection) AdditionalCmd(d bson.D) *DropCollection { + if dc == nil { + dc = new(DropCollection) + } + + dc.additionalCmd = d + return dc +} + // Session sets the session for this operation. func (dc *DropCollection) Session(session *session.Client) *DropCollection { if dc == nil { diff --git a/x/mongo/driver/operation/drop_indexes.go b/x/mongo/driver/operation/drop_indexes.go index d7fff7b3b1..ae9ecc6bb1 100644 --- a/x/mongo/driver/operation/drop_indexes.go +++ b/x/mongo/driver/operation/drop_indexes.go @@ -12,6 +12,7 @@ import ( "fmt" "time" + "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/event" "go.mongodb.org/mongo-driver/v2/internal/driverutil" "go.mongodb.org/mongo-driver/v2/mongo/writeconcern" @@ -40,6 +41,7 @@ type DropIndexes struct { serverAPI *driver.ServerAPIOptions timeout *time.Duration rawData *bool + additionalCmd bson.D } // DropIndexesResult represents a dropIndexes result returned by the server. @@ -124,6 +126,13 @@ func (di *DropIndexes) command(dst []byte, desc description.SelectedServer) ([]b if di.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) { dst = bsoncore.AppendBooleanElement(dst, "rawData", *di.rawData) } + if len(di.additionalCmd) > 0 { + doc, err := bson.Marshal(di.additionalCmd) + if err != nil { + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) + } + dst = append(dst, doc[4:len(doc)-1]...) + } return dst, nil } @@ -289,3 +298,13 @@ func (di *DropIndexes) RawData(rawData bool) *DropIndexes { di.rawData = &rawData return di } + +// AdditionalCmd sets additional command fields to be attached. +func (di *DropIndexes) AdditionalCmd(d bson.D) *DropIndexes { + if di == nil { + di = new(DropIndexes) + } + + di.additionalCmd = d + return di +} diff --git a/x/mongo/driver/operation/find.go b/x/mongo/driver/operation/find.go index 23599814bb..68289519df 100644 --- a/x/mongo/driver/operation/find.go +++ b/x/mongo/driver/operation/find.go @@ -9,8 +9,10 @@ package operation import ( "context" "errors" + "fmt" "time" + "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/event" "go.mongodb.org/mongo-driver/v2/internal/driverutil" "go.mongodb.org/mongo-driver/v2/internal/logger" @@ -64,6 +66,7 @@ type Find struct { serverAPI *driver.ServerAPIOptions timeout *time.Duration rawData *bool + additionalCmd bson.D logger *logger.Logger omitMaxTimeMS bool } @@ -200,6 +203,13 @@ func (f *Find) command(dst []byte, desc description.SelectedServer) ([]byte, err if f.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) { dst = bsoncore.AppendBooleanElement(dst, "rawData", *f.rawData) } + if len(f.additionalCmd) > 0 { + doc, err := bson.Marshal(f.additionalCmd) + if err != nil { + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) + } + dst = append(dst, doc[4:len(doc)-1]...) + } return dst, nil } @@ -606,6 +616,16 @@ func (f *Find) RawData(rawData bool) *Find { return f } +// AdditionalCmd sets additional command fields to be attached. +func (f *Find) AdditionalCmd(d bson.D) *Find { + if f == nil { + f = new(Find) + } + + f.additionalCmd = d + return f +} + // OmitMaxTimeMS omits the automatically-calculated "maxTimeMS" from the // command. func (f *Find) OmitMaxTimeMS(omit bool) *Find { diff --git a/x/mongo/driver/operation/find_and_modify.go b/x/mongo/driver/operation/find_and_modify.go index c3bd217285..809bdd794f 100644 --- a/x/mongo/driver/operation/find_and_modify.go +++ b/x/mongo/driver/operation/find_and_modify.go @@ -216,7 +216,7 @@ func (fam *FindAndModify) command(dst []byte, desc description.SelectedServer) ( if len(fam.additionalCmd) > 0 { doc, err := bson.Marshal(fam.additionalCmd) if err != nil { - return nil, err + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) } dst = append(dst, doc[4:len(doc)-1]...) } diff --git a/x/mongo/driver/operation/list_indexes.go b/x/mongo/driver/operation/list_indexes.go index 1eea4e6409..2bced33544 100644 --- a/x/mongo/driver/operation/list_indexes.go +++ b/x/mongo/driver/operation/list_indexes.go @@ -9,8 +9,10 @@ package operation import ( "context" "errors" + "fmt" "time" + "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/event" "go.mongodb.org/mongo-driver/v2/internal/driverutil" "go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore" @@ -37,6 +39,7 @@ type ListIndexes struct { serverAPI *driver.ServerAPIOptions timeout *time.Duration rawData *bool + additionalCmd bson.D result driver.CursorResponse } @@ -106,6 +109,13 @@ func (li *ListIndexes) command(dst []byte, desc description.SelectedServer) ([]b if li.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) { dst = bsoncore.AppendBooleanElement(dst, "rawData", *li.rawData) } + if len(li.additionalCmd) > 0 { + doc, err := bson.Marshal(li.additionalCmd) + if err != nil { + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) + } + dst = append(dst, doc[4:len(doc)-1]...) + } return dst, nil } @@ -272,3 +282,13 @@ func (li *ListIndexes) RawData(rawData bool) *ListIndexes { li.rawData = &rawData return li } + +// AdditionalCmd sets additional command fields to be attached. +func (li *ListIndexes) AdditionalCmd(d bson.D) *ListIndexes { + if li == nil { + li = new(ListIndexes) + } + + li.additionalCmd = d + return li +} diff --git a/x/mongo/driver/operation/update.go b/x/mongo/driver/operation/update.go index f06f729ef3..c134c6de22 100644 --- a/x/mongo/driver/operation/update.go +++ b/x/mongo/driver/operation/update.go @@ -214,7 +214,7 @@ func (u *Update) command(dst []byte, desc description.SelectedServer) ([]byte, e if len(u.additionalCmd) > 0 { doc, err := bson.Marshal(u.additionalCmd) if err != nil { - return nil, err + return nil, fmt.Errorf("error marshaling additional command fields: %w", err) } dst = append(dst, doc[4:len(doc)-1]...) } diff --git a/x/mongo/driver/xoptions/options.go b/x/mongo/driver/xoptions/options.go index c3a0938ebe..accffb9121 100644 --- a/x/mongo/driver/xoptions/options.go +++ b/x/mongo/driver/xoptions/options.go @@ -141,6 +141,15 @@ func SetInternalCountOptions(a *options.CountOptionsBuilder, key string, option opts.Internal = optionsutil.WithValue(opts.Internal, key, b) return nil }) + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.CountOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) default: return fmt.Errorf("unsupported option: %q", key) } @@ -162,6 +171,15 @@ func SetInternalCreateIndexesOptions(a *options.CreateIndexesOptionsBuilder, key opts.Internal = optionsutil.WithValue(opts.Internal, key, b) return nil }) + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.CreateIndexesOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) default: return fmt.Errorf("unsupported option: %q", key) } @@ -183,6 +201,15 @@ func SetInternalDeleteOneOptions(a *options.DeleteOneOptionsBuilder, key string, opts.Internal = optionsutil.WithValue(opts.Internal, key, b) return nil }) + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.DeleteOneOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) default: return fmt.Errorf("unsupported option: %q", key) } @@ -204,6 +231,15 @@ func SetInternalDeleteManyOptions(a *options.DeleteManyOptionsBuilder, key strin opts.Internal = optionsutil.WithValue(opts.Internal, key, b) return nil }) + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.DeleteManyOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) default: return fmt.Errorf("unsupported option: %q", key) } @@ -225,6 +261,36 @@ func SetInternalDistinctOptions(a *options.DistinctOptionsBuilder, key string, o opts.Internal = optionsutil.WithValue(opts.Internal, key, b) return nil }) + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.DistinctOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) + default: + return fmt.Errorf("unsupported option: %q", key) + } + return nil +} + +// SetInternalDropCollectionOptions sets internal options for DropCollectionOptions. +func SetInternalDropCollectionOptions(a *options.DropCollectionOptionsBuilder, key string, option any) error { + typeErrFunc := func(t string) error { + return fmt.Errorf("unexpected type for %q: %T is not %s", key, option, t) + } + switch key { + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.DropCollectionOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) default: return fmt.Errorf("unsupported option: %q", key) } @@ -246,6 +312,15 @@ func SetInternalDropIndexesOptions(a *options.DropIndexesOptionsBuilder, key str opts.Internal = optionsutil.WithValue(opts.Internal, key, b) return nil }) + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.DropIndexesOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) default: return fmt.Errorf("unsupported option: %q", key) } @@ -267,6 +342,15 @@ func SetInternalEstimatedDocumentCountOptions(a *options.EstimatedDocumentCountO opts.Internal = optionsutil.WithValue(opts.Internal, key, b) return nil }) + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.EstimatedDocumentCountOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) default: return fmt.Errorf("unsupported option: %q", key) } @@ -288,6 +372,15 @@ func SetInternalFindOptions(a *options.FindOptionsBuilder, key string, option an opts.Internal = optionsutil.WithValue(opts.Internal, key, b) return nil }) + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.FindOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) default: return fmt.Errorf("unsupported option: %q", key) } @@ -309,6 +402,15 @@ func SetInternalFindOneOptions(a *options.FindOneOptionsBuilder, key string, opt opts.Internal = optionsutil.WithValue(opts.Internal, key, b) return nil }) + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.FindOneOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) default: return fmt.Errorf("unsupported option: %q", key) } @@ -492,6 +594,15 @@ func SetInternalListIndexesOptions(a *options.ListIndexesOptionsBuilder, key str opts.Internal = optionsutil.WithValue(opts.Internal, key, b) return nil }) + case "addCommandFields": + d, ok := option.(bson.D) + if !ok { + return typeErrFunc("bson.D") + } + a.Opts = append(a.Opts, func(opts *options.ListIndexesOptions) error { + opts.Internal = optionsutil.WithValue(opts.Internal, key, d) + return nil + }) default: return fmt.Errorf("unsupported option: %q", key) } diff --git a/x/mongo/driver/xoptions/options_test.go b/x/mongo/driver/xoptions/options_test.go index 4a4e1a371e..85bd027bbb 100644 --- a/x/mongo/driver/xoptions/options_test.go +++ b/x/mongo/driver/xoptions/options_test.go @@ -10,6 +10,8 @@ import ( "fmt" "testing" + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/internal/mongoutil" "go.mongodb.org/mongo-driver/v2/internal/optionsutil" "go.mongodb.org/mongo-driver/v2/internal/require" "go.mongodb.org/mongo-driver/v2/mongo/options" @@ -87,3 +89,165 @@ func TestSetInternalClientOptions(t *testing.T) { require.EqualError(t, err, "unsupported option: \"unsupported\"") }) } + +// TestSetInternalAddCommandFields verifies that each collection- and +// index-level options setter that supports the "addCommandFields" key stores +// the provided bson.D in the resulting Internal options. +func TestSetInternalAddCommandFields(t *testing.T) { + t.Parallel() + + want := bson.D{{Key: "collectionUUID", Value: "00000000-0000-0000-0000-000000000000"}} + + // Each case sets "addCommandFields" via a specific setter and returns the + // Internal options that the setter populated. + cases := []struct { + name string + set func(t *testing.T) optionsutil.Options + }{ + {"BulkWriteOptions", func(t *testing.T) optionsutil.Options { + o := options.BulkWrite() + require.NoError(t, SetInternalBulkWriteOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.BulkWriteOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"CountOptions", func(t *testing.T) optionsutil.Options { + o := options.Count() + require.NoError(t, SetInternalCountOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.CountOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"CreateIndexesOptions", func(t *testing.T) optionsutil.Options { + o := options.CreateIndexes() + require.NoError(t, SetInternalCreateIndexesOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.CreateIndexesOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"DeleteOneOptions", func(t *testing.T) optionsutil.Options { + o := options.DeleteOne() + require.NoError(t, SetInternalDeleteOneOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.DeleteOneOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"DeleteManyOptions", func(t *testing.T) optionsutil.Options { + o := options.DeleteMany() + require.NoError(t, SetInternalDeleteManyOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.DeleteManyOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"DistinctOptions", func(t *testing.T) optionsutil.Options { + o := options.Distinct() + require.NoError(t, SetInternalDistinctOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.DistinctOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"DropCollectionOptions", func(t *testing.T) optionsutil.Options { + o := options.DropCollection() + require.NoError(t, SetInternalDropCollectionOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.DropCollectionOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"DropIndexesOptions", func(t *testing.T) optionsutil.Options { + o := options.DropIndexes() + require.NoError(t, SetInternalDropIndexesOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.DropIndexesOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"EstimatedDocumentCountOptions", func(t *testing.T) optionsutil.Options { + o := options.EstimatedDocumentCount() + require.NoError(t, SetInternalEstimatedDocumentCountOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.EstimatedDocumentCountOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"FindOptions", func(t *testing.T) optionsutil.Options { + o := options.Find() + require.NoError(t, SetInternalFindOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.FindOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"FindOneOptions", func(t *testing.T) optionsutil.Options { + o := options.FindOne() + require.NoError(t, SetInternalFindOneOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.FindOneOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"FindOneAndReplaceOptions", func(t *testing.T) optionsutil.Options { + o := options.FindOneAndReplace() + require.NoError(t, SetInternalFindOneAndReplaceOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.FindOneAndReplaceOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"FindOneAndUpdateOptions", func(t *testing.T) optionsutil.Options { + o := options.FindOneAndUpdate() + require.NoError(t, SetInternalFindOneAndUpdateOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.FindOneAndUpdateOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"InsertOneOptions", func(t *testing.T) optionsutil.Options { + o := options.InsertOne() + require.NoError(t, SetInternalInsertOneOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.InsertOneOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"InsertManyOptions", func(t *testing.T) optionsutil.Options { + o := options.InsertMany() + require.NoError(t, SetInternalInsertManyOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.InsertManyOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"ListIndexesOptions", func(t *testing.T) optionsutil.Options { + o := options.ListIndexes() + require.NoError(t, SetInternalListIndexesOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.ListIndexesOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"ReplaceOptions", func(t *testing.T) optionsutil.Options { + o := options.Replace() + require.NoError(t, SetInternalReplaceOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.ReplaceOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"UpdateOneOptions", func(t *testing.T) optionsutil.Options { + o := options.UpdateOne() + require.NoError(t, SetInternalUpdateOneOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.UpdateOneOptions](o) + require.NoError(t, err) + return args.Internal + }}, + {"UpdateManyOptions", func(t *testing.T) optionsutil.Options { + o := options.UpdateMany() + require.NoError(t, SetInternalUpdateManyOptions(o, "addCommandFields", want)) + args, err := mongoutil.NewOptions[options.UpdateManyOptions](o) + require.NoError(t, err) + return args.Internal + }}, + } + + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + internal := tc.set(t) + got, ok := optionsutil.Value(internal, "addCommandFields").(bson.D) + require.True(t, ok, "expected addCommandFields to be a bson.D") + require.Equal(t, want, got) + }) + } +} From d667702a9e70b5cefd84d58a6711fedcfb409fa3 Mon Sep 17 00:00:00 2001 From: Felipe Gasper Date: Fri, 3 Jul 2026 09:32:41 -0400 Subject: [PATCH 2/3] add integration tests --- internal/integration/collection_test.go | 313 ++++++++++++++++++++++++ 1 file changed, 313 insertions(+) diff --git a/internal/integration/collection_test.go b/internal/integration/collection_test.go index 513ffc4cc7..170ab44781 100644 --- a/internal/integration/collection_test.go +++ b/internal/integration/collection_test.go @@ -2409,6 +2409,319 @@ func TestBypassEmptyTsReplacement(t *testing.T) { }) } +// TestAddCommandFields verifies that the "addCommandFields" internal option +// injects arbitrary top-level fields into the command sent to the server for +// each collection- and index-level method that supports it. Each subtest +// injects a "comment" field (accepted by these commands on 4.4+) and asserts +// that it appears in the started command event. +func TestAddCommandFields(t *testing.T) { + mt := mtest.New(t, mtest.NewOptions().MinServerVersion("4.4")) + + const wantComment = "addCommandFieldsTest" + + marshalValue := func(val interface{}) bson.RawValue { + t.Helper() + + valType, data, err := bson.MarshalValue(val) + require.NoError(t, err, "MarshalValue error: %v", err) + return bson.RawValue{ + Type: valType, + Value: data, + } + } + + added := bson.D{{"comment", wantComment}} + empty := bson.RawValue{} + set := marshalValue(wantComment) + + mt.Run("delete one", func(mt *mtest.T) { + newOpts := func() *options.DeleteOneOptionsBuilder { + opts := options.DeleteOne() + err := xoptions.SetInternalDeleteOneOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.DeleteOneOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + _, err := mt.Coll.DeleteOne(context.Background(), bson.D{{"x", 1}}, tc.opts) + require.NoError(mt, err, "DeleteOne error: %v", err) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) + mt.Run("delete many", func(mt *mtest.T) { + newOpts := func() *options.DeleteManyOptionsBuilder { + opts := options.DeleteMany() + err := xoptions.SetInternalDeleteManyOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.DeleteManyOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + _, err := mt.Coll.DeleteMany(context.Background(), bson.D{{"x", 1}}, tc.opts) + require.NoError(mt, err, "DeleteMany error: %v", err) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) + mt.Run("find", func(mt *mtest.T) { + newOpts := func() *options.FindOptionsBuilder { + opts := options.Find() + err := xoptions.SetInternalFindOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.FindOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + initCollection(mt, mt.Coll) + mt.ClearEvents() + + cursor, err := mt.Coll.Find(context.Background(), bson.D{}, tc.opts) + require.NoError(mt, err, "Find error: %v", err) + _ = cursor.Close(context.Background()) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) + mt.Run("find one", func(mt *mtest.T) { + newOpts := func() *options.FindOneOptionsBuilder { + opts := options.FindOne() + err := xoptions.SetInternalFindOneOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.FindOneOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + initCollection(mt, mt.Coll) + mt.ClearEvents() + + err := mt.Coll.FindOne(context.Background(), bson.D{}, tc.opts).Err() + require.NoError(mt, err, "FindOne error: %v", err) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) + mt.Run("distinct", func(mt *mtest.T) { + newOpts := func() *options.DistinctOptionsBuilder { + opts := options.Distinct() + err := xoptions.SetInternalDistinctOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.DistinctOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + initCollection(mt, mt.Coll) + mt.ClearEvents() + + err := mt.Coll.Distinct(context.Background(), "x", bson.D{}, tc.opts).Err() + require.NoError(mt, err, "Distinct error: %v", err) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) + mt.Run("count documents", func(mt *mtest.T) { + newOpts := func() *options.CountOptionsBuilder { + opts := options.Count() + err := xoptions.SetInternalCountOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.CountOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + initCollection(mt, mt.Coll) + mt.ClearEvents() + + _, err := mt.Coll.CountDocuments(context.Background(), bson.D{}, tc.opts) + require.NoError(mt, err, "CountDocuments error: %v", err) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) + mt.Run("estimated document count", func(mt *mtest.T) { + newOpts := func() *options.EstimatedDocumentCountOptionsBuilder { + opts := options.EstimatedDocumentCount() + err := xoptions.SetInternalEstimatedDocumentCountOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.EstimatedDocumentCountOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + initCollection(mt, mt.Coll) + mt.ClearEvents() + + _, err := mt.Coll.EstimatedDocumentCount(context.Background(), tc.opts) + require.NoError(mt, err, "EstimatedDocumentCount error: %v", err) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) + mt.Run("list indexes", func(mt *mtest.T) { + newOpts := func() *options.ListIndexesOptionsBuilder { + opts := options.ListIndexes() + err := xoptions.SetInternalListIndexesOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.ListIndexesOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + initCollection(mt, mt.Coll) + mt.ClearEvents() + + cursor, err := mt.Coll.Indexes().List(context.Background(), tc.opts) + require.NoError(mt, err, "List error: %v", err) + _ = cursor.Close(context.Background()) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) + mt.Run("create indexes", func(mt *mtest.T) { + newOpts := func() *options.CreateIndexesOptionsBuilder { + opts := options.CreateIndexes() + err := xoptions.SetInternalCreateIndexesOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.CreateIndexesOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + model := mongo.IndexModel{Keys: bson.D{{"x", 1}}} + _, err := mt.Coll.Indexes().CreateOne(context.Background(), model, tc.opts) + require.NoError(mt, err, "CreateOne error: %v", err) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) + mt.Run("drop indexes", func(mt *mtest.T) { + newOpts := func() *options.DropIndexesOptionsBuilder { + opts := options.DropIndexes() + err := xoptions.SetInternalDropIndexesOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.DropIndexesOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + name, err := mt.Coll.Indexes().CreateOne(context.Background(), + mongo.IndexModel{Keys: bson.D{{"x", 1}}}) + require.NoError(mt, err, "CreateOne error: %v", err) + mt.ClearEvents() + + err = mt.Coll.Indexes().DropOne(context.Background(), name, tc.opts) + require.NoError(mt, err, "DropOne error: %v", err) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) +} + func initCollection(tb testing.TB, coll *mongo.Collection) { tb.Helper() From 6fb45b82b92505c8ca2d8f552aa9799c5f53c565 Mon Sep 17 00:00:00 2001 From: Felipe Gasper Date: Fri, 3 Jul 2026 10:08:11 -0400 Subject: [PATCH 3/3] Copilot notes --- internal/integration/collection_test.go | 34 +++++++++++++++++++++++++ mongo/collection.go | 13 +++++++++- x/mongo/driver/operation/aggregate.go | 20 --------------- 3 files changed, 46 insertions(+), 21 deletions(-) diff --git a/internal/integration/collection_test.go b/internal/integration/collection_test.go index 170ab44781..fa5b50c74d 100644 --- a/internal/integration/collection_test.go +++ b/internal/integration/collection_test.go @@ -2632,6 +2632,40 @@ func TestAddCommandFields(t *testing.T) { }) } }) + mt.Run("drop", func(mt *mtest.T) { + newOpts := func() *options.DropCollectionOptionsBuilder { + opts := options.DropCollection() + err := xoptions.SetInternalDropCollectionOptions(opts, "addCommandFields", added) + require.NoError(mt, err, "unexpected error: %v", err) + return opts + } + + testCases := []struct { + name string + opts *options.DropCollectionOptionsBuilder + expected bson.RawValue + }{ + {"empty", nil, empty}, + {"set", newOpts(), set}, + } + for _, tc := range testCases { + mt.Run(tc.name, func(mt *mtest.T) { + initCollection(mt, mt.Coll) + mt.ClearEvents() + + var err error + if tc.opts == nil { + err = mt.Coll.Drop(context.Background()) + } else { + err = mt.Coll.Drop(context.Background(), tc.opts) + } + require.NoError(mt, err, "Drop error: %v", err) + evt := mt.GetStartedEvent() + val := evt.Command.Lookup("comment") + assert.Equal(mt, tc.expected, val, "expected comment to be %s", tc.expected.String()) + }) + } + }) mt.Run("list indexes", func(mt *mtest.T) { newOpts := func() *options.ListIndexesOptionsBuilder { opts := options.ListIndexes() diff --git a/mongo/collection.go b/mongo/collection.go index b85cdfa46f..6829a270cf 100644 --- a/mongo/collection.go +++ b/mongo/collection.go @@ -1217,8 +1217,19 @@ func (coll *Collection) CountDocuments(ctx context.Context, filter any, if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { op = op.RawData(rawData) } + // CountDocuments runs an aggregate command, which exposes arbitrary + // top-level command fields through CustomOptions rather than a dedicated + // additionalCmd builder. if additionalCmd, ok := optionsutil.Value(args.Internal, "addCommandFields").(bson.D); ok { - op = op.AdditionalCmd(additionalCmd) + customOptions := make(map[string]bsoncore.Value) + for _, elem := range additionalCmd { + elemValueBSON, err := marshalValue(elem.Value, coll.bsonOpts, coll.registry) + if err != nil { + return 0, err + } + customOptions[elem.Key] = elemValueBSON + } + op.CustomOptions(customOptions) } err = op.Execute(ctx) diff --git a/x/mongo/driver/operation/aggregate.go b/x/mongo/driver/operation/aggregate.go index 94baa8ed3d..22dbb8087f 100644 --- a/x/mongo/driver/operation/aggregate.go +++ b/x/mongo/driver/operation/aggregate.go @@ -9,10 +9,8 @@ package operation import ( "context" "errors" - "fmt" "time" - "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/event" "go.mongodb.org/mongo-driver/v2/internal/driverutil" "go.mongodb.org/mongo-driver/v2/mongo/readconcern" @@ -55,7 +53,6 @@ type Aggregate struct { timeout *time.Duration omitMaxTimeMS bool rawData *bool - additionalCmd bson.D result driver.CursorResponse } @@ -164,13 +161,6 @@ func (a *Aggregate) command(dst []byte, desc description.SelectedServer) ([]byte if a.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) { dst = bsoncore.AppendBooleanElement(dst, "rawData", *a.rawData) } - if len(a.additionalCmd) > 0 { - doc, err := bson.Marshal(a.additionalCmd) - if err != nil { - return nil, fmt.Errorf("error marshaling additional command fields: %w", err) - } - dst = append(dst, doc[4:len(doc)-1]...) - } for optionName, optionValue := range a.customOptions { dst = bsoncore.AppendValueElement(dst, optionName, optionValue) } @@ -475,13 +465,3 @@ func (a *Aggregate) RawData(rawData bool) *Aggregate { a.rawData = &rawData return a } - -// AdditionalCmd sets additional command fields to be attached. -func (a *Aggregate) AdditionalCmd(d bson.D) *Aggregate { - if a == nil { - a = new(Aggregate) - } - - a.additionalCmd = d - return a -}