add count

Author: qingyang-hu

mongo/collection.go

@@ -1127,6 +1127,9 @@ func (coll *Collection) CountDocuments(ctx context.Context, filter interface{},
 		}
 		op.Hint(hintVal)
 	}
+	if args.RawData != nil {
+		op = op.RawData(*args.RawData)
+	}
 	retry := driver.RetryNone
 	if coll.client.retryReads {
 		retry = driver.RetryOncePerCommand
@@ -1208,6 +1211,9 @@ func (coll *Collection) EstimatedDocumentCount(
 		}
 		op = op.Comment(comment)
 	}
+	if args.RawData != nil {
+		op = op.RawData(*args.RawData)
+	}
 
 	retry := driver.RetryNone
 	if coll.client.retryReads {

mongo/options/countoptions.go

@@ -16,6 +16,7 @@ type CountOptions struct {
 	Hint      interface{}
 	Limit     *int64
 	Skip      *int64
+	RawData   *bool
 }
 
 // CountOptionsBuilder contains options to configure count operations. Each
@@ -99,3 +100,15 @@ func (co *CountOptionsBuilder) SetSkip(i int64) *CountOptionsBuilder {
 
 	return co
 }
+
+// SetRawData sets the value for the RawData field. If true, it allows the CRUD operations to access timeseries
+// collections on the bucket-level. This option is only valid for MongoDB versions >= 9.0. The default value is false.
+func (co *CountOptionsBuilder) SetRawData(rawData bool) *CountOptionsBuilder {
+	co.Opts = append(co.Opts, func(opts *CountOptions) error {
+		opts.RawData = &rawData
+
+		return nil
+	})
+
+	return co
+}

mongo/options/estimatedcountoptions.go

@@ -12,6 +12,7 @@ package options
 // See corresponding setter methods for documentation.
 type EstimatedDocumentCountOptions struct {
 	Comment interface{}
+	RawData *bool
 }
 
 // EstimatedDocumentCountOptionsBuilder contains options to estimate document
@@ -44,3 +45,15 @@ func (eco *EstimatedDocumentCountOptionsBuilder) SetComment(comment interface{})
 
 	return eco
 }
+
+// SetRawData sets the value for the RawData field. If true, it allows the CRUD operations to access timeseries
+// collections on the bucket-level. This option is only valid for MongoDB versions >= 9.0. The default value is false.
+func (eco *EstimatedDocumentCountOptionsBuilder) SetRawData(rawData bool) *EstimatedDocumentCountOptionsBuilder {
+	eco.Opts = append(eco.Opts, func(opts *EstimatedDocumentCountOptions) error {
+		opts.RawData = &rawData
+
+		return nil
+	})
+
+	return eco
+}

x/mongo/driver/operation/count.go

@@ -41,6 +41,7 @@ type Count struct {
 	result         CountResult
 	serverAPI      *driver.ServerAPIOptions
 	timeout        *time.Duration
+	rawData        *bool
 }
 
 // CountResult represents a count result returned by the server.
@@ -139,14 +140,18 @@ func (c *Count) Execute(ctx context.Context) error {
 	return err
 }
 
-func (c *Count) command(dst []byte, _ description.SelectedServer) ([]byte, error) {
+func (c *Count) command(dst []byte, desc description.SelectedServer) ([]byte, error) {
 	dst = bsoncore.AppendStringElement(dst, "count", c.collection)
 	if c.query != nil {
 		dst = bsoncore.AppendDocumentElement(dst, "query", c.query)
 	}
 	if c.comment.Type != bsoncore.Type(0) {
 		dst = bsoncore.AppendValueElement(dst, "comment", c.comment)
 	}
+	// Set rawData for 8.2+ servers.
+	if c.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) {
+		dst = bsoncore.AppendBooleanElement(dst, "rawData", *c.rawData)
+	}
 	return dst, nil
 }
 
@@ -310,3 +315,13 @@ func (c *Count) Authenticator(authenticator driver.Authenticator) *Count {
 	c.authenticator = authenticator
 	return c
 }
+
+// RawData sets the rawData to access timeseries data in the compressed format.
+func (c *Count) RawData(rawData bool) *Count {
+	if c == nil {
+		c = new(Count)
+	}
+
+	c.rawData = &rawData
+	return c
+}