GODRIVER-1150 generate Count operation

Change-Id: Idbfeaa9aa27ceb03511b765cd2977ab78765b7b1

Author: iwysiu

mongo/collection.go

@@ -19,7 +19,6 @@ import (
 	"go.mongodb.org/mongo-driver/mongo/readconcern"
 	"go.mongodb.org/mongo-driver/mongo/readpref"
 	"go.mongodb.org/mongo-driver/mongo/writeconcern"
-	"go.mongodb.org/mongo-driver/x/bsonx"
 	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 	"go.mongodb.org/mongo-driver/x/mongo/driver"
 	"go.mongodb.org/mongo-driver/x/mongo/driver/description"
@@ -787,42 +786,39 @@ func (coll *Collection) EstimatedDocumentCount(ctx context.Context,
 
 	sess := sessionFromContext(ctx)
 
+	if sess == nil && coll.client.topology.SessionPool != nil {
+		sess, err := session.NewClientSession(coll.client.topology.SessionPool, coll.client.id, session.Implicit)
+		if err != nil {
+			return 0, err
+		}
+		defer sess.EndSession()
+	}
+
 	err := coll.client.validSession(sess)
 	if err != nil {
 		return 0, err
 	}
 
 	rc := coll.readConcern
-	if sess != nil && (sess.TransactionInProgress()) {
+	if sess.TransactionRunning() {
 		rc = nil
 	}
 
-	oldns := coll.namespace()
-	cmd := command.Count{
-		NS:          command.Namespace{DB: oldns.DB, Collection: oldns.Collection},
-		Query:       bsonx.Doc{},
-		ReadPref:    coll.readPreference,
-		ReadConcern: rc,
-		Session:     sess,
-		Clock:       coll.client.clock,
-	}
+	selector := makePinnedSelector(sess, coll.readSelector)
 
-	countOpts := options.Count()
-	if len(opts) >= 1 {
-		countOpts = countOpts.SetMaxTime(*opts[len(opts)-1].MaxTime)
+	op := operation.NewCount().Session(sess).ClusterClock(coll.client.clock).
+		Database(coll.db.name).Collection(coll.name).CommandMonitor(coll.client.monitor).
+		Deployment(coll.client.topology).ReadConcern(rc).ReadPreference(coll.readPreference).
+		ServerSelector(selector)
+
+	co := options.MergeEstimatedDocumentCountOptions(opts...)
+	if co.MaxTime != nil {
+		op = op.MaxTimeMS(int64(*co.MaxTime / time.Millisecond))
 	}
 
-	count, err := driverlegacy.Count(
-		ctx, cmd,
-		coll.client.topology,
-		coll.readSelector,
-		coll.client.id,
-		coll.client.topology.SessionPool,
-		coll.registry,
-		countOpts,
-	)
+	err = op.Execute(ctx)
 
-	return count, replaceErrors(err)
+	return op.Result().N, replaceErrors(err)
 }
 
 // Distinct finds the distinct values for a specified field across a single

mongo/session.go

@@ -182,12 +182,7 @@ func (s *sessionImpl) AbortTransaction(ctx context.Context) error {
 		return s.clientSession.AbortTransaction()
 	}
 
-	var selector description.ServerSelectorFunc = func(t description.Topology, svrs []description.Server) ([]description.Server, error) {
-		if s.clientSession.PinnedServer != nil {
-			return s.clientSession.PinnedServer.SelectServer(t, svrs)
-		}
-		return description.WriteSelector().SelectServer(t, svrs)
-	}
+	selector := makePinnedSelector(s.clientSession, description.WriteSelector())
 
 	s.clientSession.Aborting = true
 	err = operation.NewAbortTransaction().Session(s.clientSession).ClusterClock(s.client.clock).Database("admin").

x/mongo/driver/operation/count.go

@@ -0,0 +1,25 @@
+version = 0
+name = "Count"
+documentation = "Performs a count operation"
+
+[properties]
+enabled = ["read concern", "read preference"]
+
+[command]
+name = "count"
+parameter = "collection"
+
+[request.maxTimeMS]
+type = "int64"
+documentation = "MaxTimeMS specifies the maximum amount of time to allow the query to run."
+
+[request.query]
+type = "document"
+documentation = "Query determines what results are returned from find."
+
+[response]
+name = "CountResult"
+
+[response.field.n]
+type = "int64"
+documentation = "The number of documents found"

x/mongo/driver/operation/count.toml

@@ -11,3 +11,4 @@ package operation
 //go:generate operationgen drop_database.toml operation drop_database.go
 //go:generate operationgen commit_transaction.toml operation commit_transaction.go
 //go:generate operationgen abort_transaction.toml operation abort_transaction.go
+//go:generate operationgen count.toml operation count.go