SERVER-73083 Support findAndModify with remove: true on a timeseries collection

1 files changed, 27 insertions, 14 deletions

diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp
index 082563180bc..83c475dfd50 100644
--- a/src/mongo/db/query/plan_executor_impl.cpp
+++ b/src/mongo/db/query/plan_executor_impl.cpp
@@ -637,27 +637,40 @@ long long PlanExecutorImpl::executeDelete() {
     }
 
     // If the collection exists, the delete plan may either have a delete stage at the root, or
-    // (for findAndModify) a projection stage wrapping a delete stage.
-    switch (_root->stageType()) {
-        case StageType::STAGE_PROJECTION_DEFAULT:
-        case StageType::STAGE_PROJECTION_COVERED:
-        case StageType::STAGE_PROJECTION_SIMPLE: {
-            invariant(_root->getChildren().size() == 1U);
-            invariant(StageType::STAGE_DELETE == _root->child()->stageType());
-            const SpecificStats* stats = _root->child()->getSpecificStats();
-            return static_cast<const DeleteStats*>(stats)->docsDeleted;
+    // (for findAndModify) a projection stage wrapping a delete / TS_MODIFY stage.
+    const auto deleteStage = [&] {
+        switch (_root->stageType()) {
+            case StageType::STAGE_PROJECTION_DEFAULT:
+            case StageType::STAGE_PROJECTION_COVERED:
+            case StageType::STAGE_PROJECTION_SIMPLE: {
+                tassert(7308302,
+                        "Unexpected number of children: {}"_format(_root->getChildren().size()),
+                        _root->getChildren().size() == 1U);
+                auto childStage = _root->child().get();
+                tassert(7308303,
+                        "Unexpected child stage type: {}"_format(childStage->stageType()),
+                        StageType::STAGE_DELETE == childStage->stageType() ||
+                            StageType::STAGE_TIMESERIES_MODIFY == childStage->stageType());
+                return childStage;
+            }
+            default:
+                return _root.get();
         }
+    }();
+    switch (deleteStage->stageType()) {
         case StageType::STAGE_TIMESERIES_MODIFY: {
             const auto* tsModifyStats =
-                static_cast<const TimeseriesModifyStats*>(_root->getSpecificStats());
+                static_cast<const TimeseriesModifyStats*>(deleteStage->getSpecificStats());
             return tsModifyStats->nMeasurementsModified;
         }
-        default: {
-            invariant(StageType::STAGE_DELETE == _root->stageType() ||
-                      StageType::STAGE_BATCHED_DELETE == _root->stageType());
-            const auto* deleteStats = static_cast<const DeleteStats*>(_root->getSpecificStats());
+        case StageType::STAGE_DELETE:
+        case StageType::STAGE_BATCHED_DELETE: {
+            const auto* deleteStats =
+                static_cast<const DeleteStats*>(deleteStage->getSpecificStats());
             return deleteStats->docsDeleted;
         }
+        default:
+            MONGO_UNREACHABLE_TASSERT(7308306);
     }
 }