GODRIVER-1149 migrate countDocuments to new aggregate

Change-Id: Ib378c7ec182993975cf769996c9b8865ce2a6f17

Author: iwysiu

mongo/collection.go

@@ -24,7 +24,6 @@ import (
 	"go.mongodb.org/mongo-driver/x/mongo/driver/description"
 	"go.mongodb.org/mongo-driver/x/mongo/driver/operation"
 	"go.mongodb.org/mongo-driver/x/mongo/driver/session"
-	"go.mongodb.org/mongo-driver/x/mongo/driverlegacy"
 	"go.mongodb.org/mongo-driver/x/network/command"
 )
 
@@ -747,9 +746,14 @@ func (coll *Collection) CountDocuments(ctx context.Context, filter interface{},
 	}
 
 	sess := sessionFromContext(ctx)
-
-	err = coll.client.validSession(sess)
-	if err != nil {
+	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()
+	}
+	if err = coll.client.validSession(sess); err != nil {
 		return 0, err
 	}
 
@@ -758,27 +762,46 @@ func (coll *Collection) CountDocuments(ctx context.Context, filter interface{},
 		rc = nil
 	}
 
-	oldns := coll.namespace()
-	cmd := command.CountDocuments{
-		NS:          command.Namespace{DB: oldns.DB, Collection: oldns.Collection},
-		Pipeline:    pipelineArr,
-		ReadPref:    coll.readPreference,
-		ReadConcern: rc,
-		Session:     sess,
-		Clock:       coll.client.clock,
+	selector := makePinnedSelector(sess, coll.readSelector)
+
+	op := operation.NewAggregate(pipelineArr).Session(sess).ReadConcern(rc).ReadPreference(coll.readPreference).
+		CommandMonitor(coll.client.monitor).ServerSelector(selector).ClusterClock(coll.client.clock).Database(coll.db.name).
+		Collection(coll.name).Deployment(coll.client.topology)
+	if countOpts.Collation != nil {
+		op.Collation(bsoncore.Document(countOpts.Collation.ToDocument()))
+	}
+	if countOpts.MaxTime != nil {
+		op.MaxTimeMS(int64(*countOpts.MaxTime / time.Millisecond))
+	}
+	if countOpts.Hint != nil {
+		hintVal, err := transformValue(coll.registry, countOpts.Hint)
+		if err != nil {
+			return 0, err
+		}
+		op.Hint(hintVal)
+	}
+
+	err = op.Execute(ctx)
+	if err != nil {
+		return 0, replaceErrors(err)
+	}
+
+	batch := op.ResultCursorResponse().FirstBatch
+	if batch == nil {
+		return 0, errors.New("invalid response from server, no 'firstBatch' field")
 	}
 
-	count, err := driverlegacy.CountDocuments(
-		ctx, cmd,
-		coll.client.topology,
-		coll.readSelector,
-		coll.client.id,
-		coll.client.topology.SessionPool,
-		coll.registry,
-		countOpts,
-	)
+	docs, err := batch.Documents()
+	if err != nil || len(docs) == 0 {
+		return 0, nil
+	}
+
+	val, ok := docs[0].Lookup("n").AsInt64OK()
+	if !ok {
+		return 0, errors.New("invalid response from server, no 'n' field")
+	}
 
-	return count, replaceErrors(err)
+	return val, nil
 }
 
 // EstimatedDocumentCount gets an estimate of the count of documents in a collection using collection metadata.

mongo/mongo.go

@@ -348,31 +348,41 @@ func transformValue(registry *bsoncodec.Registry, val interface{}) (bsoncore.Val
 }
 
 // Build the aggregation pipeline for the CountDocument command.
-func countDocumentsAggregatePipeline(registry *bsoncodec.Registry, filter interface{}, opts *options.CountOptions) (bsonx.Arr, error) {
-	pipeline := bsonx.Arr{}
-	filterDoc, err := transformDocument(registry, filter)
-
+func countDocumentsAggregatePipeline(registry *bsoncodec.Registry, filter interface{}, opts *options.CountOptions) (bsoncore.Document, error) {
+	filterDoc, err := transformBsoncoreDocument(registry, filter)
 	if err != nil {
 		return nil, err
 	}
-	pipeline = append(pipeline, bsonx.Document(bsonx.Doc{{"$match", bsonx.Document(filterDoc)}}))
 
+	aidx, arr := bsoncore.AppendArrayStart(nil)
+	didx, arr := bsoncore.AppendDocumentElementStart(arr, strconv.Itoa(0))
+	arr = bsoncore.AppendDocumentElement(arr, "$match", filterDoc)
+	arr, _ = bsoncore.AppendDocumentEnd(arr, didx)
+
+	index := 1
 	if opts != nil {
 		if opts.Skip != nil {
-			pipeline = append(pipeline, bsonx.Document(bsonx.Doc{{"$skip", bsonx.Int64(*opts.Skip)}}))
+			didx, arr = bsoncore.AppendDocumentElementStart(arr, strconv.Itoa(index))
+			arr = bsoncore.AppendInt64Element(arr, "$skip", *opts.Skip)
+			arr, _ = bsoncore.AppendDocumentEnd(arr, didx)
+			index++
 		}
 		if opts.Limit != nil {
-			pipeline = append(pipeline, bsonx.Document(bsonx.Doc{{"$limit", bsonx.Int64(*opts.Limit)}}))
+			didx, arr = bsoncore.AppendDocumentElementStart(arr, strconv.Itoa(index))
+			arr = bsoncore.AppendInt64Element(arr, "$limit", *opts.Limit)
+			arr, _ = bsoncore.AppendDocumentEnd(arr, didx)
+			index++
 		}
 	}
 
-	pipeline = append(pipeline, bsonx.Document(bsonx.Doc{
-		{"$group", bsonx.Document(bsonx.Doc{
-			{"_id", bsonx.Int32(1)},
-			{"n", bsonx.Document(bsonx.Doc{{"$sum", bsonx.Int32(1)}})},
-		})},
-	},
-	))
+	didx, arr = bsoncore.AppendDocumentElementStart(arr, strconv.Itoa(index))
+	iidx, arr := bsoncore.AppendDocumentElementStart(arr, "$group")
+	arr = bsoncore.AppendInt32Element(arr, "_id", 1)
+	iiidx, arr := bsoncore.AppendDocumentElementStart(arr, "n")
+	arr = bsoncore.AppendInt32Element(arr, "$sum", 1)
+	arr, _ = bsoncore.AppendDocumentEnd(arr, iiidx)
+	arr, _ = bsoncore.AppendDocumentEnd(arr, iidx)
+	arr, _ = bsoncore.AppendDocumentEnd(arr, didx)
 
-	return pipeline, nil
+	return bsoncore.AppendArrayEnd(arr, aidx)
 }