SERVER-35479 Make $collStats count behavior have "standard" error code and consistent behavior across topologies

4 files changed, 32 insertions, 24 deletions

diff --git a/src/mongo/db/pipeline/document_source_coll_stats.cpp b/src/mongo/db/pipeline/document_source_coll_stats.cpp
index 72c504a6093..1c2d295b27b 100644
--- a/src/mongo/db/pipeline/document_source_coll_stats.cpp
+++ b/src/mongo/db/pipeline/document_source_coll_stats.cpp
@@ -131,34 +131,23 @@ DocumentSource::GetNextResult DocumentSourceCollStats::doGetNext() {
     if (_collStatsSpec.hasField("storageStats")) {
         // If the storageStats field exists, it must have been validated as an object when parsing.
         BSONObjBuilder storageBuilder(builder.subobjStart("storageStats"));
-        Status status = pExpCtx->mongoProcessInterface->appendStorageStats(
-            pExpCtx->opCtx, pExpCtx->ns, _collStatsSpec["storageStats"].Obj(), &storageBuilder);
+        uassertStatusOKWithContext(
+            pExpCtx->mongoProcessInterface->appendStorageStats(
+                pExpCtx->opCtx, pExpCtx->ns, _collStatsSpec["storageStats"].Obj(), &storageBuilder),
+            "Unable to retrieve storageStats in $collStats stage");
         storageBuilder.doneFast();
-        if (!status.isOK()) {
-            uasserted(40280,
-                      str::stream() << "Unable to retrieve storageStats in $collStats stage: "
-                                    << status.reason());
-        }
     }
 
     if (_collStatsSpec.hasField("count")) {
-        Status status = pExpCtx->mongoProcessInterface->appendRecordCount(
-            pExpCtx->opCtx, pExpCtx->ns, &builder);
-        if (!status.isOK()) {
-            uasserted(40481,
-                      str::stream()
-                          << "Unable to retrieve count in $collStats stage: " << status.reason());
-        }
+        uassertStatusOKWithContext(pExpCtx->mongoProcessInterface->appendRecordCount(
+                                       pExpCtx->opCtx, pExpCtx->ns, &builder),
+                                   "Unable to retrieve count in $collStats stage");
     }
 
     if (_collStatsSpec.hasField("queryExecStats")) {
-        Status status = pExpCtx->mongoProcessInterface->appendQueryExecStats(
-            pExpCtx->opCtx, pExpCtx->ns, &builder);
-        if (!status.isOK()) {
-            uasserted(31142,
-                      str::stream() << "Unable to retrieve queryExecStats in $collStats stage: "
-                                    << status.reason());
-        }
+        uassertStatusOKWithContext(pExpCtx->mongoProcessInterface->appendQueryExecStats(
+                                       pExpCtx->opCtx, pExpCtx->ns, &builder),
+                                   "Unable to retrieve queryExecStats in $collStats stage");
     }
 
     return {Document(builder.obj())};
diff --git a/src/mongo/db/pipeline/document_source_coll_stats.h b/src/mongo/db/pipeline/document_source_coll_stats.h
index 6e97cf562c6..4d0be020523 100644
--- a/src/mongo/db/pipeline/document_source_coll_stats.h
+++ b/src/mongo/db/pipeline/document_source_coll_stats.h
@@ -45,11 +45,14 @@ public:
     public:
         static std::unique_ptr<LiteParsed> parse(const NamespaceString& nss,
                                                  const BSONElement& spec) {
-            return std::make_unique<LiteParsed>(spec.fieldName(), nss);
+            bool hasCount = spec.isABSONObj() && spec.embeddedObject().hasField("count");
+            return std::make_unique<LiteParsed>(spec.fieldName(), nss, hasCount);
         }
 
-        explicit LiteParsed(std::string parseTimeName, NamespaceString nss)
-            : LiteParsedDocumentSource(std::move(parseTimeName)), _nss(std::move(nss)) {}
+        explicit LiteParsed(std::string parseTimeName, NamespaceString nss, bool hasCount)
+            : LiteParsedDocumentSource(std::move(parseTimeName)),
+              _nss(std::move(nss)),
+              _hasCount(hasCount) {}
 
         bool isCollStats() const final {
             return true;
@@ -68,8 +71,13 @@ public:
             return true;
         }
 
+        bool isCollStatsWithCount() const final {
+            return _hasCount;
+        }
+
     private:
         const NamespaceString _nss;
+        const bool _hasCount;
     };
 
     DocumentSourceCollStats(const boost::intrusive_ptr<ExpressionContext>& pExpCtx)
diff --git a/src/mongo/db/pipeline/lite_parsed_document_source.h b/src/mongo/db/pipeline/lite_parsed_document_source.h
index 95c28f49510..38c01ab02c3 100644
--- a/src/mongo/db/pipeline/lite_parsed_document_source.h
+++ b/src/mongo/db/pipeline/lite_parsed_document_source.h
@@ -105,6 +105,10 @@ public:
         return false;
     }
 
+    virtual bool isCollStatsWithCount() const {
+        return false;
+    }
+
     /**
      * Returns true if this is a $changeStream stage.
      */
diff --git a/src/mongo/db/pipeline/lite_parsed_pipeline.h b/src/mongo/db/pipeline/lite_parsed_pipeline.h
index d80a470c9b0..0af48f8b74f 100644
--- a/src/mongo/db/pipeline/lite_parsed_pipeline.h
+++ b/src/mongo/db/pipeline/lite_parsed_pipeline.h
@@ -97,6 +97,13 @@ public:
     }
 
     /**
+     * Returns true if the pipeline begins with a $collStats stage with the count option.
+     */
+    bool startsWithCollStatsWithCount() const {
+        return startsWithCollStats() && _stageSpecs.front()->isCollStatsWithCount();
+    }
+
+    /**
      * Returns true if the pipeline has a $changeStream stage.
      */
     bool hasChangeStream() const {