SERVER-29620 Add shorthand syntax for retrieving the record count for a collection using the agg stage

8 files changed, 138 insertions, 1 deletions

diff --git a/jstests/aggregation/sources/collStats/count.js b/jstests/aggregation/sources/collStats/count.js
new file mode 100644
index 00000000000..4072a9d56fc
--- /dev/null
+++ b/jstests/aggregation/sources/collStats/count.js
@@ -0,0 +1,65 @@
+// Test that count within a $collStats stage returns the correct number of documents.
+// @tags: [assumes_no_implicit_collection_creation_after_drop]
+(function() {
+    "use strict";
+
+    load("jstests/aggregation/extras/utils.js");  // For "assertErrorCode".
+
+    let testDB = db.getSiblingDB("aggregation_count_db");
+    let coll = testDB.aggregation_count;
+    coll.drop();
+
+    let nDocs = 1000;
+    for (var i = 0; i < nDocs; i++) {
+        assert.writeOK(coll.insert({a: i}));
+    }
+
+    // Test that $collStats must be first stage.
+    let pipeline = [{$match: {}}, {$collStats: {}}];
+    assertErrorCode(coll, pipeline, ErrorCodes.BadValue);
+
+    // Test that an error is returned if count is not an object.
+    pipeline = [{$collStats: {count: 1}}];
+    assertErrorCode(coll, pipeline, 40480, "count spec must be an object");
+    pipeline = [{$collStats: {count: "1"}}];
+    assertErrorCode(coll, pipeline, 40480, "count spec must be an object");
+
+    // Test the accuracy of the record count as a standalone option.
+    pipeline = [{$collStats: {count: {}}}];
+    let result = coll.aggregate(pipeline).next();
+    assert.eq(nDocs, result.count);
+
+    // Test the record count alongside latencyStats and storageStats.
+    pipeline = [{$collStats: {count: {}, latencyStats: {}}}];
+    result = coll.aggregate(pipeline).next();
+    assert.eq(nDocs, result.count);
+    assert(result.hasOwnProperty("latencyStats"));
+    assert(result.latencyStats.hasOwnProperty("reads"));
+    assert(result.latencyStats.hasOwnProperty("writes"));
+    assert(result.latencyStats.hasOwnProperty("commands"));
+
+    pipeline = [{$collStats: {count: {}, latencyStats: {}, storageStats: {}}}];
+    result = coll.aggregate(pipeline).next();
+    assert.eq(nDocs, result.count);
+    assert(result.hasOwnProperty("latencyStats"));
+    assert(result.latencyStats.hasOwnProperty("reads"));
+    assert(result.latencyStats.hasOwnProperty("writes"));
+    assert(result.latencyStats.hasOwnProperty("commands"));
+    assert(result.hasOwnProperty("storageStats"));
+    assert.eq(nDocs, result.storageStats.count);
+
+    // Test the record count against an empty collection.
+    assert.writeOK(coll.remove({}));
+    pipeline = [{$collStats: {count: {}}}];
+    result = coll.aggregate(pipeline).next();
+    assert.eq(0, result.count);
+
+    // Assert when there is no collection.
+    coll.drop();
+    assertErrorCode(coll, pipeline, 40481);
+
+    // Assert when there is no database.
+    assert.commandWorked(testDB.dropDatabase());
+    assertErrorCode(coll, pipeline, 40481);
+
+}());
diff --git a/jstests/core/views/views_coll_stats.js b/jstests/core/views/views_coll_stats.js
index 5e7d9587f43..bae2aa4e41f 100644
--- a/jstests/core/views/views_coll_stats.js
+++ b/jstests/core/views/views_coll_stats.js
@@ -64,4 +64,20 @@
         viewsDB.runCommand(
             {aggregate: "a", pipeline: [{$collStats: {storageStats: {}}}], cursor: {}}),
         ErrorCodes.CommandNotSupportedOnView);
+    clear();
+
+    // Assert that attempting to retrieve collection record count on an identity views fails.
+    makeView("a", "b");
+    assert.commandFailedWithCode(
+        viewsDB.runCommand({aggregate: "a", pipeline: [{$collStats: {count: {}}}], cursor: {}}),
+        ErrorCodes.CommandNotSupportedOnView);
+    clear();
+
+    // Assert that attempting to retrieve collection record count on a non-identity view fails.
+    makeView("a", "b", [{$match: {a: 0}}]);
+    assert.commandFailedWithCode(
+        viewsDB.runCommand({aggregate: "a", pipeline: [{$collStats: {count: {}}}], cursor: {}}),
+        ErrorCodes.CommandNotSupportedOnView);
+    clear();
+
 }());
diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
index fab65e25173..994f17151b3 100644
--- a/src/mongo/db/pipeline/document_source.h
+++ b/src/mongo/db/pipeline/document_source.h
@@ -599,6 +599,12 @@ public:
                                           BSONObjBuilder* builder) const = 0;
 
         /**
+         * Appends the record count for collection "nss" to "builder".
+         */
+        virtual Status appendRecordCount(const NamespaceString& nss,
+                                         BSONObjBuilder* builder) const = 0;
+
+        /**
          * Gets the collection options for the collection given by 'nss'.
          */
         virtual BSONObj getCollectionOptions(const NamespaceString& nss) = 0;
diff --git a/src/mongo/db/pipeline/document_source_coll_stats.cpp b/src/mongo/db/pipeline/document_source_coll_stats.cpp
index bdf2ee35b5c..c4516ff2cf5 100644
--- a/src/mongo/db/pipeline/document_source_coll_stats.cpp
+++ b/src/mongo/db/pipeline/document_source_coll_stats.cpp
@@ -77,6 +77,12 @@ intrusive_ptr<DocumentSource> DocumentSourceCollStats::createFromBson(
                                   << " of type "
                                   << typeName(elem.type()),
                     elem.type() == BSONType::Object);
+        } else if ("count" == fieldName) {
+            uassert(40480,
+                    str::stream() << "count argument must be an object, but got " << elem
+                                  << " of type "
+                                  << typeName(elem.type()),
+                    elem.type() == BSONType::Object);
         } else {
             uasserted(40168, str::stream() << "unrecognized option to $collStats: " << fieldName);
         }
@@ -119,6 +125,15 @@ DocumentSource::GetNextResult DocumentSourceCollStats::getNext() {
         }
     }
 
+    if (_collStatsSpec.hasField("count")) {
+        Status status = _mongod->appendRecordCount(pExpCtx->ns, &builder);
+        if (!status.isOK()) {
+            uasserted(40481,
+                      str::stream() << "Unable to retrieve count in $collStats stage: "
+                                    << status.reason());
+        }
+    }
+
     return {Document(builder.obj())};
 }
 
diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp
index 9c0eb1ec9bd..a27f403430b 100644
--- a/src/mongo/db/pipeline/pipeline_d.cpp
+++ b/src/mongo/db/pipeline/pipeline_d.cpp
@@ -142,6 +142,10 @@ public:
         return appendCollectionStorageStats(_ctx->opCtx, nss, param, builder);
     }
 
+    Status appendRecordCount(const NamespaceString& nss, BSONObjBuilder* builder) const final {
+        return appendCollectionRecordCount(_ctx->opCtx, nss, builder);
+    }
+
     BSONObj getCollectionOptions(const NamespaceString& nss) final {
         const auto infos =
             _client.getCollectionInfos(nss.db().toString(), BSON("name" << nss.coll()));
diff --git a/src/mongo/db/pipeline/stub_mongod_interface.h b/src/mongo/db/pipeline/stub_mongod_interface.h
index e75e98e91f8..6f20ea79451 100644
--- a/src/mongo/db/pipeline/stub_mongod_interface.h
+++ b/src/mongo/db/pipeline/stub_mongod_interface.h
@@ -75,6 +75,10 @@ public:
         MONGO_UNREACHABLE;
     }
 
+    Status appendRecordCount(const NamespaceString& nss, BSONObjBuilder* builder) const override {
+        MONGO_UNREACHABLE;
+    }
+
     BSONObj getCollectionOptions(const NamespaceString& nss) override {
         MONGO_UNREACHABLE;
     }
diff --git a/src/mongo/db/stats/storage_stats.cpp b/src/mongo/db/stats/storage_stats.cpp
index fca73b4a473..4d45dc02b4b 100644
--- a/src/mongo/db/stats/storage_stats.cpp
+++ b/src/mongo/db/stats/storage_stats.cpp
@@ -108,4 +108,24 @@ Status appendCollectionStorageStats(OperationContext* opCtx,
 
     return Status::OK();
 }
+
+Status appendCollectionRecordCount(OperationContext* opCtx,
+                                   const NamespaceString& nss,
+                                   BSONObjBuilder* result) {
+    AutoGetCollectionForReadCommand ctx(opCtx, nss);
+    if (!ctx.getDb()) {
+        return {ErrorCodes::BadValue,
+                str::stream() << "Database [" << nss.db().toString() << "] not found."};
+    }
+
+    Collection* collection = ctx.getCollection();
+    if (!collection) {
+        return {ErrorCodes::BadValue,
+                str::stream() << "Collection [" << nss.toString() << "] not found."};
+    }
+
+    result->appendNumber("count", static_cast<long long>(collection->numRecords(opCtx)));
+
+    return Status::OK();
+}
 }  // namespace mongo
diff --git a/src/mongo/db/stats/storage_stats.h b/src/mongo/db/stats/storage_stats.h
index 0f62229c80a..f3a3db99b9b 100644
--- a/src/mongo/db/stats/storage_stats.h
+++ b/src/mongo/db/stats/storage_stats.h
@@ -38,10 +38,17 @@ namespace mongo {
 
 /**
  * Appends to 'builder' storage related statistics for the collection represented by 'nss'.
- * Used by both the collStats command and the $collStats aggregation stage.
  */
 Status appendCollectionStorageStats(OperationContext* opCtx,
                                     const NamespaceString& nss,
                                     const BSONObj& param,
                                     BSONObjBuilder* builder);
+
+/**
+ * Appends the collection record count to 'builder' for the collection represented by 'nss'.
+ */
+Status appendCollectionRecordCount(OperationContext* opCtx,
+                                   const NamespaceString& nss,
+                                   BSONObjBuilder* builder);
+
 };  // namespace mongo