From 969385a01ba83be0a3f69f7c83b5fa99141a443f Mon Sep 17 00:00:00 2001 From: Matt Dale <9760375+matthewdale@users.noreply.github.com> Date: Fri, 10 Jul 2026 23:36:39 -0700 Subject: [PATCH 1/2] GODRIVER-3892 Move count to the mongo package. --- mongo/collection.go | 29 ++- mongo/op_count.go | 150 +++++++++++++ x/mongo/driver/operation/count.go | 352 ------------------------------ 3 files changed, 170 insertions(+), 361 deletions(-) create mode 100644 mongo/op_count.go delete mode 100644 x/mongo/driver/operation/count.go diff --git a/mongo/collection.go b/mongo/collection.go index 3f1df6dbd..ecece89c6 100644 --- a/mongo/collection.go +++ b/mongo/collection.go @@ -1299,23 +1299,34 @@ func (coll *Collection) EstimatedDocumentCount( maxAdaptiveRetries := coll.client.effectiveAdaptiveRetries(coll.client.retryReads) selector := makeReadPrefSelector(sess, coll.readSelector, coll.client.localThreshold) - op := operation.NewCount().Session(sess).ClusterClock(coll.client.clock). - Database(coll.db.name).Collection(coll.name).CommandMonitor(coll.client.monitor). - Deployment(coll.client.deployment).ReadConcern(rc).ReadPreference(coll.readPreference). - Retry(retry).MaxAdaptiveRetries(maxAdaptiveRetries). - EnableOverloadRetargeting(coll.client.enableOverloadRetargeting). - ServerSelector(selector).Crypt(coll.client.cryptFLE).ServerAPI(coll.client.serverAPI). - Timeout(coll.client.timeout).Authenticator(coll.client.authenticator) + op := countOp{ + session: sess, + clock: coll.client.clock, + database: coll.db.name, + collection: coll.name, + monitor: coll.client.monitor, + deployment: coll.client.deployment, + readConcern: rc, + readPreference: coll.readPreference, + retry: &retry, + maxAdaptiveRetries: maxAdaptiveRetries, + enableOverloadRetargeting: coll.client.enableOverloadRetargeting, + selector: selector, + crypt: coll.client.cryptFLE, + serverAPI: coll.client.serverAPI, + timeout: coll.client.timeout, + authenticator: coll.client.authenticator, + } if args.Comment != nil { comment, err := marshalValue(args.Comment, coll.bsonOpts, coll.registry) if err != nil { return 0, err } - op = op.Comment(comment) + op.comment = comment } if rawData, ok := optionsutil.Value(args.Internal, "rawData").(bool); ok { - op = op.RawData(rawData) + op.rawData = &rawData } err = op.Execute(ctx) diff --git a/mongo/op_count.go b/mongo/op_count.go new file mode 100644 index 000000000..1fc47e0eb --- /dev/null +++ b/mongo/op_count.go @@ -0,0 +1,150 @@ +// Copyright (C) MongoDB, Inc. 2019-present. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may +// not use this file except in compliance with the License. You may obtain +// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 + +package mongo + +import ( + "context" + "errors" + "fmt" + "time" + + "go.mongodb.org/mongo-driver/v2/event" + "go.mongodb.org/mongo-driver/v2/internal/driverutil" + "go.mongodb.org/mongo-driver/v2/mongo/readconcern" + "go.mongodb.org/mongo-driver/v2/mongo/readpref" + "go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore" + "go.mongodb.org/mongo-driver/v2/x/mongo/driver" + "go.mongodb.org/mongo-driver/v2/x/mongo/driver/description" + "go.mongodb.org/mongo-driver/v2/x/mongo/driver/session" +) + +// countOp performs a count operation. +type countOp struct { + authenticator driver.Authenticator + session *session.Client + clock *session.ClusterClock + collection string + comment bsoncore.Value + monitor *event.CommandMonitor + crypt driver.Crypt + database string + deployment driver.Deployment + readConcern *readconcern.ReadConcern + readPreference *readpref.ReadPref + selector description.ServerSelector + retry *driver.RetryMode + maxAdaptiveRetries uint + enableOverloadRetargeting bool + result countResult + serverAPI *driver.ServerAPIOptions + timeout *time.Duration + rawData *bool +} + +// countResult represents a count result returned by the server. +type countResult struct { + // The number of documents found + N int64 +} + +func buildCountResult(response bsoncore.Document) (countResult, error) { + elements, err := response.Elements() + if err != nil { + return countResult{}, err + } + cr := countResult{} + for _, element := range elements { + switch element.Key() { + case "n": // for count using original command + var ok bool + cr.N, ok = element.Value().AsInt64OK() + if !ok { + return cr, fmt.Errorf("response field 'n' is type int64, but received BSON type %s", + element.Value().Type) + } + case "cursor": // for count using aggregate with $collStats + firstBatch, err := element.Value().Document().LookupErr("firstBatch") + if err != nil { + return cr, err + } + + // get count value from first batch + val := firstBatch.Array().Index(0) + count, err := val.Document().LookupErr("n") + if err != nil { + return cr, err + } + + // use count as Int64 for result + var ok bool + cr.N, ok = count.AsInt64OK() + if !ok { + return cr, fmt.Errorf("response field 'n' is type int64, but received BSON type %s", + element.Value().Type) + } + } + } + return cr, nil +} + +// Result returns the result of executing this operation. +func (c *countOp) Result() countResult { return c.result } + +func (c *countOp) processResponse(_ context.Context, resp bsoncore.Document, _ driver.ResponseInfo) error { + var err error + c.result, err = buildCountResult(resp) + return err +} + +// Execute runs this operations and returns an error if the operation did not execute successfully. +func (c *countOp) Execute(ctx context.Context) error { + if c.deployment == nil { + return errors.New("the Count operation must have a Deployment set before Execute can be called") + } + + err := driver.Operation{ + CommandFn: c.command, + ProcessResponseFn: c.processResponse, + RetryMode: c.retry, + MaxAdaptiveRetries: c.maxAdaptiveRetries, + EnableOverloadRetargeting: c.enableOverloadRetargeting, + Type: driver.Read, + Client: c.session, + Clock: c.clock, + CommandMonitor: c.monitor, + Crypt: c.crypt, + Database: c.database, + Deployment: c.deployment, + ReadConcern: c.readConcern, + ReadPreference: c.readPreference, + Selector: c.selector, + ServerAPI: c.serverAPI, + Timeout: c.timeout, + Name: driverutil.CountOp, + Authenticator: c.authenticator, + }.Execute(ctx) + // Swallow error if NamespaceNotFound(26) is returned from aggregate on non-existent namespace + if err != nil { + dErr, ok := err.(driver.Error) + if ok && dErr.Code == 26 { + err = nil + } + } + return err +} + +func (c *countOp) command(dst []byte, desc description.SelectedServer) ([]byte, error) { + dst = bsoncore.AppendStringElement(dst, "count", c.collection) + 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 +} diff --git a/x/mongo/driver/operation/count.go b/x/mongo/driver/operation/count.go deleted file mode 100644 index 69e585ebf..000000000 --- a/x/mongo/driver/operation/count.go +++ /dev/null @@ -1,352 +0,0 @@ -// Copyright (C) MongoDB, Inc. 2019-present. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may -// not use this file except in compliance with the License. You may obtain -// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 - -package operation - -import ( - "context" - "errors" - "fmt" - "time" - - "go.mongodb.org/mongo-driver/v2/event" - "go.mongodb.org/mongo-driver/v2/internal/driverutil" - "go.mongodb.org/mongo-driver/v2/mongo/readconcern" - "go.mongodb.org/mongo-driver/v2/mongo/readpref" - "go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore" - "go.mongodb.org/mongo-driver/v2/x/mongo/driver" - "go.mongodb.org/mongo-driver/v2/x/mongo/driver/description" - "go.mongodb.org/mongo-driver/v2/x/mongo/driver/session" -) - -// Count represents a count operation. -type Count struct { - authenticator driver.Authenticator - query bsoncore.Document - session *session.Client - clock *session.ClusterClock - collection string - comment bsoncore.Value - monitor *event.CommandMonitor - crypt driver.Crypt - database string - deployment driver.Deployment - readConcern *readconcern.ReadConcern - readPreference *readpref.ReadPref - selector description.ServerSelector - retry *driver.RetryMode - maxAdaptiveRetries uint - enableOverloadRetargeting bool - result CountResult - serverAPI *driver.ServerAPIOptions - timeout *time.Duration - rawData *bool -} - -// CountResult represents a count result returned by the server. -type CountResult struct { - // The number of documents found - N int64 -} - -func buildCountResult(response bsoncore.Document) (CountResult, error) { - elements, err := response.Elements() - if err != nil { - return CountResult{}, err - } - cr := CountResult{} - for _, element := range elements { - switch element.Key() { - case "n": // for count using original command - var ok bool - cr.N, ok = element.Value().AsInt64OK() - if !ok { - return cr, fmt.Errorf("response field 'n' is type int64, but received BSON type %s", - element.Value().Type) - } - case "cursor": // for count using aggregate with $collStats - firstBatch, err := element.Value().Document().LookupErr("firstBatch") - if err != nil { - return cr, err - } - - // get count value from first batch - val := firstBatch.Array().Index(0) - count, err := val.Document().LookupErr("n") - if err != nil { - return cr, err - } - - // use count as Int64 for result - var ok bool - cr.N, ok = count.AsInt64OK() - if !ok { - return cr, fmt.Errorf("response field 'n' is type int64, but received BSON type %s", - element.Value().Type) - } - } - } - return cr, nil -} - -// NewCount constructs and returns a new Count. -func NewCount() *Count { - return &Count{} -} - -// Result returns the result of executing this operation. -func (c *Count) Result() CountResult { return c.result } - -func (c *Count) processResponse(_ context.Context, resp bsoncore.Document, _ driver.ResponseInfo) error { - var err error - c.result, err = buildCountResult(resp) - return err -} - -// Execute runs this operations and returns an error if the operation did not execute successfully. -func (c *Count) Execute(ctx context.Context) error { - if c.deployment == nil { - return errors.New("the Count operation must have a Deployment set before Execute can be called") - } - - err := driver.Operation{ - CommandFn: c.command, - ProcessResponseFn: c.processResponse, - RetryMode: c.retry, - MaxAdaptiveRetries: c.maxAdaptiveRetries, - EnableOverloadRetargeting: c.enableOverloadRetargeting, - Type: driver.Read, - Client: c.session, - Clock: c.clock, - CommandMonitor: c.monitor, - Crypt: c.crypt, - Database: c.database, - Deployment: c.deployment, - ReadConcern: c.readConcern, - ReadPreference: c.readPreference, - Selector: c.selector, - ServerAPI: c.serverAPI, - Timeout: c.timeout, - Name: driverutil.CountOp, - Authenticator: c.authenticator, - }.Execute(ctx) - // Swallow error if NamespaceNotFound(26) is returned from aggregate on non-existent namespace - if err != nil { - dErr, ok := err.(driver.Error) - if ok && dErr.Code == 26 { - err = nil - } - } - return err -} - -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 -} - -// Query determines what results are returned from find. -func (c *Count) Query(query bsoncore.Document) *Count { - if c == nil { - c = new(Count) - } - - c.query = query - return c -} - -// Session sets the session for this operation. -func (c *Count) Session(session *session.Client) *Count { - if c == nil { - c = new(Count) - } - - c.session = session - return c -} - -// ClusterClock sets the cluster clock for this operation. -func (c *Count) ClusterClock(clock *session.ClusterClock) *Count { - if c == nil { - c = new(Count) - } - - c.clock = clock - return c -} - -// Collection sets the collection that this command will run against. -func (c *Count) Collection(collection string) *Count { - if c == nil { - c = new(Count) - } - - c.collection = collection - return c -} - -// Comment sets a value to help trace an operation. -func (c *Count) Comment(comment bsoncore.Value) *Count { - if c == nil { - c = new(Count) - } - - c.comment = comment - return c -} - -// CommandMonitor sets the monitor to use for APM events. -func (c *Count) CommandMonitor(monitor *event.CommandMonitor) *Count { - if c == nil { - c = new(Count) - } - - c.monitor = monitor - return c -} - -// Crypt sets the Crypt object to use for automatic encryption and decryption. -func (c *Count) Crypt(crypt driver.Crypt) *Count { - if c == nil { - c = new(Count) - } - - c.crypt = crypt - return c -} - -// Database sets the database to run this operation against. -func (c *Count) Database(database string) *Count { - if c == nil { - c = new(Count) - } - - c.database = database - return c -} - -// Deployment sets the deployment to use for this operation. -func (c *Count) Deployment(deployment driver.Deployment) *Count { - if c == nil { - c = new(Count) - } - - c.deployment = deployment - return c -} - -// ReadConcern specifies the read concern for this operation. -func (c *Count) ReadConcern(readConcern *readconcern.ReadConcern) *Count { - if c == nil { - c = new(Count) - } - - c.readConcern = readConcern - return c -} - -// ReadPreference set the read preference used with this operation. -func (c *Count) ReadPreference(readPreference *readpref.ReadPref) *Count { - if c == nil { - c = new(Count) - } - - c.readPreference = readPreference - return c -} - -// ServerSelector sets the selector used to retrieve a server. -func (c *Count) ServerSelector(selector description.ServerSelector) *Count { - if c == nil { - c = new(Count) - } - - c.selector = selector - return c -} - -// Retry enables retryable mode for this operation. Retries are handled automatically in driver.Operation.Execute based -// on how the operation is set. -func (c *Count) Retry(retry driver.RetryMode) *Count { - if c == nil { - c = new(Count) - } - - c.retry = &retry - return c -} - -// MaxAdaptiveRetries specifies the maximum number of times the driver should retry operations -// that fail with a server side overload error. -func (c *Count) MaxAdaptiveRetries(maxAdaptiveRetries uint) *Count { - if c == nil { - c = new(Count) - } - - c.maxAdaptiveRetries = maxAdaptiveRetries - return c -} - -// EnableOverloadRetargeting specifies whether the driver adds the previously failed server's address -// to the list of deprioritized server addresses -func (c *Count) EnableOverloadRetargeting(enabled bool) *Count { - if c == nil { - c = new(Count) - } - - c.enableOverloadRetargeting = enabled - return c -} - -// ServerAPI sets the server API version for this operation. -func (c *Count) ServerAPI(serverAPI *driver.ServerAPIOptions) *Count { - if c == nil { - c = new(Count) - } - - c.serverAPI = serverAPI - return c -} - -// Timeout sets the timeout for this operation. -func (c *Count) Timeout(timeout *time.Duration) *Count { - if c == nil { - c = new(Count) - } - - c.timeout = timeout - return c -} - -// Authenticator sets the authenticator to use for this operation. -func (c *Count) Authenticator(authenticator driver.Authenticator) *Count { - if c == nil { - c = new(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 -} From 908dbd07b42563be4052c89a2c5d501b6c2d4008 Mon Sep 17 00:00:00 2001 From: Matt Dale <9760375+matthewdale@users.noreply.github.com> Date: Tue, 14 Jul 2026 10:13:40 -0700 Subject: [PATCH 2/2] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- mongo/op_count.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mongo/op_count.go b/mongo/op_count.go index 1fc47e0eb..6f7a2ad6a 100644 --- a/mongo/op_count.go +++ b/mongo/op_count.go @@ -100,10 +100,10 @@ func (c *countOp) processResponse(_ context.Context, resp bsoncore.Document, _ d return err } -// Execute runs this operations and returns an error if the operation did not execute successfully. +// Execute runs this operation and returns an error if the operation did not execute successfully. func (c *countOp) Execute(ctx context.Context) error { if c.deployment == nil { - return errors.New("the Count operation must have a Deployment set before Execute can be called") + return errors.New("the count operation must have a Deployment set before Execute can be called") } err := driver.Operation{