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

6 files changed, 57 insertions, 1 deletions

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