Add rawData option to aggregate, count/countDocuments/estimatedDocumentCount, delete, distinct, and insert (#2121)

Author: qingyang-hu

internal/integration/unified/collection_operation_execution.go

@@ -19,6 +19,7 @@ import (
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 	"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
+	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/xoptions"
 )
 
 // This file contains helpers to execute collection operations.
@@ -75,6 +76,11 @@ func executeAggregate(ctx context.Context, operation *operation) (*operationResu
 			pipeline = bsonutil.RawToInterfaces(bsonutil.RawArrayToDocuments(val.Array())...)
 		case "let":
 			opts.SetLet(val.Document())
+		case "rawData":
+			err = xoptions.SetInternalAggregateOptions(opts, key, val.Boolean())
+			if err != nil {
+				return nil, err
+			}
 		default:
 			return nil, fmt.Errorf("unrecognized aggregate option %q", key)
 		}
@@ -202,6 +208,11 @@ func executeCountDocuments(ctx context.Context, operation *operation) (*operatio
 			return nil, fmt.Errorf("the maxTimeMS collection option is not supported")
 		case "skip":
 			opts.SetSkip(int64(val.Int32()))
+		case "rawData":
+			err = xoptions.SetInternalCountOptions(opts, key, val.Boolean())
+			if err != nil {
+				return nil, err
+			}
 		default:
 			return nil, fmt.Errorf("unrecognized countDocuments option %q", key)
 		}
@@ -433,6 +444,11 @@ func executeDeleteOne(ctx context.Context, operation *operation) (*operationResu
 			opts.SetHint(hint)
 		case "let":
 			opts.SetLet(val.Document())
+		case "rawData":
+			err = xoptions.SetInternalDeleteOneOptions(opts, key, val.Boolean())
+			if err != nil {
+				return nil, err
+			}
 		default:
 			return nil, fmt.Errorf("unrecognized deleteOne option %q", key)
 		}
@@ -487,6 +503,11 @@ func executeDeleteMany(ctx context.Context, operation *operation) (*operationRes
 			opts.SetHint(hint)
 		case "let":
 			opts.SetLet(val.Document())
+		case "rawData":
+			err = xoptions.SetInternalDeleteManyOptions(opts, key, val.Boolean())
+			if err != nil {
+				return nil, err
+			}
 		default:
 			return nil, fmt.Errorf("unrecognized deleteMany option %q", key)
 		}
@@ -545,6 +566,11 @@ func executeDistinct(ctx context.Context, operation *operation) (*operationResul
 			// ensured an analogue exists, extend "skippedTestDescriptions" to avoid
 			// this error.
 			return nil, fmt.Errorf("the maxTimeMS collection option is not supported")
+		case "rawData":
+			err = xoptions.SetInternalDistinctOptions(opts, key, val.Boolean())
+			if err != nil {
+				return nil, err
+			}
 		default:
 			return nil, fmt.Errorf("unrecognized distinct option %q", key)
 		}
@@ -690,6 +716,11 @@ func executeEstimatedDocumentCount(ctx context.Context, operation *operation) (*
 			// ensured an analogue exists, extend "skippedTestDescriptions" to avoid
 			// this error.
 			return nil, fmt.Errorf("the maxTimeMS collection option is not supported")
+		case "rawData":
+			err = xoptions.SetInternalEstimatedDocumentCountOptions(opts, key, val.Boolean())
+			if err != nil {
+				return nil, err
+			}
 		default:
 			return nil, fmt.Errorf("unrecognized estimatedDocumentCount option %q", key)
 		}
@@ -1062,6 +1093,11 @@ func executeInsertMany(ctx context.Context, operation *operation) (*operationRes
 			documents = bsonutil.RawToInterfaces(bsonutil.RawArrayToDocuments(val.Array())...)
 		case "ordered":
 			opts.SetOrdered(val.Boolean())
+		case "rawData":
+			err = xoptions.SetInternalInsertManyOptions(opts, key, val.Boolean())
+			if err != nil {
+				return nil, err
+			}
 		default:
 			return nil, fmt.Errorf("unrecognized insertMany option %q", key)
 		}
@@ -1112,6 +1148,11 @@ func executeInsertOne(ctx context.Context, operation *operation) (*operationResu
 			opts.SetBypassDocumentValidation(val.Boolean())
 		case "comment":
 			opts.SetComment(val)
+		case "rawData":
+			err = xoptions.SetInternalInsertOneOptions(opts, key, val.Boolean())
+			if err != nil {
+				return nil, err
+			}
 		default:
 			return nil, fmt.Errorf("unrecognized insertOne option %q", key)
 		}

mongo/collection.go

@@ -17,6 +17,7 @@ import (
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/internal/csfle"
 	"go.mongodb.org/mongo-driver/v2/internal/mongoutil"
+	"go.mongodb.org/mongo-driver/v2/internal/optionsutil"
 	"go.mongodb.org/mongo-driver/v2/internal/serverselector"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 	"go.mongodb.org/mongo-driver/v2/mongo/readconcern"
@@ -324,6 +325,11 @@ func (coll *Collection) insert(
 	if args.Ordered != nil {
 		op = op.Ordered(*args.Ordered)
 	}
+	if rawDataOpt := optionsutil.Value(args.Internal, "rawData"); rawDataOpt != nil {
+		if rawData, ok := rawDataOpt.(bool); ok {
+			op = op.RawData(rawData)
+		}
+	}
 	retry := driver.RetryNone
 	if coll.client.retryWrites {
 		retry = driver.RetryOncePerCommand
@@ -375,6 +381,13 @@ func (coll *Collection) InsertOne(ctx context.Context, document interface{},
 	if args.Comment != nil {
 		imOpts.SetComment(args.Comment)
 	}
+	if rawDataOpt := optionsutil.Value(args.Internal, "rawData"); rawDataOpt != nil {
+		imOpts.Opts = append(imOpts.Opts, func(opts *options.InsertManyOptions) error {
+			optionsutil.WithValue(opts.Internal, "rawData", rawDataOpt)
+
+			return nil
+		})
+	}
 	res, err := coll.insert(ctx, []interface{}{document}, imOpts)
 
 	rr, err := processWriteError(err)
@@ -534,6 +547,11 @@ func (coll *Collection) delete(
 		}
 		op = op.Let(let)
 	}
+	if rawDataOpt := optionsutil.Value(args.Internal, "rawData"); rawDataOpt != nil {
+		if rawData, ok := rawDataOpt.(bool); ok {
+			op = op.RawData(rawData)
+		}
+	}
 
 	// deleteMany cannot be retried
 	retryMode := driver.RetryNone
@@ -575,6 +593,7 @@ func (coll *Collection) DeleteOne(
 		Comment:   args.Comment,
 		Hint:      args.Hint,
 		Let:       args.Let,
+		Internal:  args.Internal,
 	}
 
 	return coll.delete(ctx, filter, true, rrOne, deleteOptions)
@@ -1036,6 +1055,11 @@ func aggregate(a aggregateParams, opts ...options.Lister[options.AggregateOption
 		}
 		op.CustomOptions(customOptions)
 	}
+	if rawDataOpt := optionsutil.Value(args.Internal, "rawData"); rawDataOpt != nil {
+		if rawData, ok := rawDataOpt.(bool); ok {
+			op = op.RawData(rawData)
+		}
+	}
 
 	retry := driver.RetryNone
 	if a.retryRead && !hasOutputStage {
@@ -1124,6 +1148,11 @@ func (coll *Collection) CountDocuments(ctx context.Context, filter interface{},
 		}
 		op.Hint(hintVal)
 	}
+	if rawDataOpt := optionsutil.Value(args.Internal, "rawData"); rawDataOpt != nil {
+		if rawData, ok := rawDataOpt.(bool); ok {
+			op = op.RawData(rawData)
+		}
+	}
 	retry := driver.RetryNone
 	if coll.client.retryReads {
 		retry = driver.RetryOncePerCommand
@@ -1205,6 +1234,11 @@ func (coll *Collection) EstimatedDocumentCount(
 		}
 		op = op.Comment(comment)
 	}
+	if rawDataOpt := optionsutil.Value(args.Internal, "rawData"); rawDataOpt != nil {
+		if rawData, ok := rawDataOpt.(bool); ok {
+			op = op.RawData(rawData)
+		}
+	}
 
 	retry := driver.RetryNone
 	if coll.client.retryReads {
@@ -1294,6 +1328,11 @@ func (coll *Collection) Distinct(
 		}
 		op.Hint(hint)
 	}
+	if rawDataOpt := optionsutil.Value(args.Internal, "rawData"); rawDataOpt != nil {
+		if rawData, ok := rawDataOpt.(bool); ok {
+			op = op.RawData(rawData)
+		}
+	}
 	retry := driver.RetryNone
 	if coll.client.retryReads {
 		retry = driver.RetryOncePerCommand

mongo/options/aggregateoptions.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/internal/optionsutil"
 )
 
 // AggregateOptions represents arguments that can be used to configure an
@@ -26,6 +27,10 @@ type AggregateOptions struct {
 	Hint                     interface{}
 	Let                      interface{}
 	Custom                   bson.M
+
+	// 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
 }
 
 // AggregateOptionsBuilder contains options to configure aggregate operations.

mongo/options/countoptions.go

@@ -6,6 +6,8 @@
 
 package options
 
+import "go.mongodb.org/mongo-driver/v2/internal/optionsutil"
+
 // CountOptions represents arguments that can be used to configure a
 // CountDocuments operation.
 //
@@ -16,6 +18,10 @@ type CountOptions struct {
 	Hint      interface{}
 	Limit     *int64
 	Skip      *int64
+
+	// 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
 }
 
 // CountOptionsBuilder contains options to configure count operations. Each

mongo/options/deleteoptions.go

@@ -6,6 +6,8 @@
 
 package options
 
+import "go.mongodb.org/mongo-driver/v2/internal/optionsutil"
+
 // DeleteOneOptions represents arguments that can be used to configure DeleteOne
 // operations.
 //
@@ -15,6 +17,10 @@ type DeleteOneOptions struct {
 	Comment   interface{}
 	Hint      interface{}
 	Let       interface{}
+
+	// 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
 }
 
 // DeleteOneOptionsBuilder contains options to configure DeleteOne operations. Each
@@ -102,6 +108,10 @@ type DeleteManyOptions struct {
 	Comment   interface{}
 	Hint      interface{}
 	Let       interface{}
+
+	// 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
 }
 
 // DeleteManyOptionsBuilder contains options to configure DeleteMany operations.

mongo/options/distinctoptions.go

@@ -6,6 +6,8 @@
 
 package options
 
+import "go.mongodb.org/mongo-driver/v2/internal/optionsutil"
+
 // DistinctOptions represents arguments that can be used to configure a Distinct
 // operation.
 //
@@ -14,6 +16,10 @@ type DistinctOptions struct {
 	Collation *Collation
 	Comment   interface{}
 	Hint      interface{}
+
+	// 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
 }
 
 // DistinctOptionsBuilder contains options to configure distinct operations. Each

mongo/options/estimatedcountoptions.go

@@ -6,12 +6,18 @@
 
 package options
 
+import "go.mongodb.org/mongo-driver/v2/internal/optionsutil"
+
 // EstimatedDocumentCountOptions represents arguments that can be used to configure
 // an EstimatedDocumentCount operation.
 //
 // See corresponding setter methods for documentation.
 type EstimatedDocumentCountOptions struct {
 	Comment interface{}
+
+	// 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
 }
 
 // EstimatedDocumentCountOptionsBuilder contains options to estimate document

mongo/options/insertoptions.go

@@ -6,13 +6,19 @@
 
 package options
 
+import "go.mongodb.org/mongo-driver/v2/internal/optionsutil"
+
 // InsertOneOptions represents arguments that can be used to configure an InsertOne
 // operation.
 //
 // See corresponding setter methods for documentation.
 type InsertOneOptions struct {
 	BypassDocumentValidation *bool
 	Comment                  interface{}
+
+	// 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
 }
 
 // InsertOneOptionsBuilder represents functional options that configure an
@@ -61,6 +67,10 @@ type InsertManyOptions struct {
 	BypassDocumentValidation *bool
 	Comment                  interface{}
 	Ordered                  *bool
+
+	// 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
 }
 
 // InsertManyOptionsBuilder contains options to configure insert operations.

x/mongo/driver/operation/aggregate.go

@@ -50,6 +50,7 @@ type Aggregate struct {
 	customOptions            map[string]bsoncore.Value
 	timeout                  *time.Duration
 	omitMaxTimeMS            bool
+	rawData                  *bool
 
 	result driver.CursorResponse
 }
@@ -159,6 +160,10 @@ func (a *Aggregate) command(dst []byte, desc description.SelectedServer) ([]byte
 	if a.let != nil {
 		dst = bsoncore.AppendDocumentElement(dst, "let", a.let)
 	}
+	// Set rawData for 8.2+ servers.
+	if a.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) {
+		dst = bsoncore.AppendBooleanElement(dst, "rawData", *a.rawData)
+	}
 	for optionName, optionValue := range a.customOptions {
 		dst = bsoncore.AppendValueElement(dst, optionName, optionValue)
 	}
@@ -431,3 +436,13 @@ func (a *Aggregate) OmitMaxTimeMS(omit bool) *Aggregate {
 	a.omitMaxTimeMS = omit
 	return a
 }
+
+// RawData sets the rawData to access timeseries data in the compressed format.
+func (a *Aggregate) RawData(rawData bool) *Aggregate {
+	if a == nil {
+		a = new(Aggregate)
+	}
+
+	a.rawData = &rawData
+	return a
+}