SERVER-53696 _internalUnpackBucket should handle missing metaField gracefully

3 files changed, 116 insertions, 22 deletions

diff --git a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
index da3bbfebcb2..c6ae7ee55c9 100644
--- a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
@@ -55,10 +55,22 @@ void BucketUnpacker::reset(Document&& bucket) {
     }
 
     _metaValue = _bucket[kBucketMetaFieldName];
-    uassert(5346511,
-            "A metadata value cannot be undefined nor missing if metaField is specified",
-            !_spec.metaField ||
-                (_metaValue.getType() != BSONType::Undefined && !_metaValue.missing()));
+    if (_spec.metaField) {
+        // The spec indicates that there should be a metadata region. Missing metadata in this case
+        // is expressed with null, so the field is expected to be present. We also disallow
+        // undefined since the undefined BSON type is deprecated.
+        uassert(5369600,
+                "The $_internalUnpackBucket stage requires metadata to be present in a bucket if "
+                "metaField parameter is provided",
+                (_metaValue.getType() != BSONType::Undefined) && !_metaValue.missing());
+    } else {
+        // If the spec indicates that the time series collection has no metadata field, then we
+        // should not find a metadata region in the underlying bucket documents.
+        uassert(5369601,
+                "The $_internalUnpackBucket stage expects buckets to have missing metadata regions "
+                "if the metaField parameter is not provided",
+                _metaValue.missing());
+    }
 
     _timeFieldIter = _bucket[kBucketDataFieldName][_spec.timeField].getDocument().fieldIterator();
 
@@ -89,7 +101,7 @@ Document BucketUnpacker::getNext() {
         measurement.addField(_spec.timeField, timeVal);
     }
 
-    if (!_metaValue.nullish()) {
+    if (_includeMetaField && !_metaValue.nullish()) {
         measurement.addField(*_spec.metaField, _metaValue);
     }
 
@@ -165,8 +177,22 @@ boost::intrusive_ptr<DocumentSource> DocumentSourceInternalUnpackBucket::createF
     auto includeTimeField = (unpackerBehavior == BucketUnpacker::Behavior::kInclude) ==
         (bucketSpec.fieldSet.find(bucketSpec.timeField) != bucketSpec.fieldSet.end());
 
+    // Check the include/exclude set to determine if measurements should be materialized with
+    // metadata.
+    auto includeMetaField = false;
+    if (bucketSpec.metaField) {
+        const auto metaFieldIt = bucketSpec.fieldSet.find(*bucketSpec.metaField);
+        auto found = metaFieldIt != bucketSpec.fieldSet.end();
+        if (found) {
+            bucketSpec.fieldSet.erase(metaFieldIt);
+        }
+        includeMetaField = (unpackerBehavior == BucketUnpacker::Behavior::kInclude) == found;
+    }
+
     return make_intrusive<DocumentSourceInternalUnpackBucket>(
-        expCtx, BucketUnpacker{std::move(bucketSpec), unpackerBehavior, includeTimeField});
+        expCtx,
+        BucketUnpacker{
+            std::move(bucketSpec), unpackerBehavior, includeTimeField, includeMetaField});
 }
 
 Value DocumentSourceInternalUnpackBucket::serialize(
diff --git a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.h b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.h
index 9df0d6d6cbd..9e4c18d273a 100644
--- a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.h
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.h
@@ -47,8 +47,7 @@ struct BucketSpec {
     // after unpacking.
     boost::optional<std::string> metaField;
 
-    // The set of fields provided to the include (or exclude) argument to the
-    // $_internalUnpackBucketStage.
+    // The set of field names in the data region that should be included or excluded.
     std::set<std::string> fieldSet;
 };
 
@@ -67,10 +66,14 @@ public:
     // set difference between all fields in the bucket and the provided fields.
     enum class Behavior { kInclude, kExclude };
 
-    BucketUnpacker(BucketSpec spec, Behavior unpackerBehavior, bool includeTimeField)
+    BucketUnpacker(BucketSpec spec,
+                   Behavior unpackerBehavior,
+                   bool includeTimeField,
+                   bool includeMetaField)
         : _spec(std::move(spec)),
           _unpackerBehavior(unpackerBehavior),
-          _includeTimeField(includeTimeField) {}
+          _includeTimeField(includeTimeField),
+          _includeMetaField(includeMetaField) {}
 
     Document getNext();
     bool hasNext() const {
@@ -108,6 +111,9 @@ private:
     // metadata value in the reset phase and use it to materialize the metadata in each measurement.
     Value _metaValue;
 
+    // A flag used to mark that a bucket's metadata value should be materialized in measurements.
+    const bool _includeMetaField;
+
     // The bucket being unpacked.
     Document _bucket;
 
diff --git a/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test.cpp b/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test.cpp
index 6378eefc4a0..b4bbc9319dd 100644
--- a/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test.cpp
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket_test.cpp
@@ -47,8 +47,7 @@ TEST_F(InternalUnpackBucketStageTest, UnpackBasicIncludeAllMeasurementFields) {
 
     auto spec = BSON("$_internalUnpackBucket"
                      << BSON("include" << BSON_ARRAY("_id"
-                                                     << "time"
-                                                     << "a"
+                                                     << "time" << kUserDefinedMetaName << "a"
                                                      << "b")
                                        << DocumentSourceInternalUnpackBucket::kTimeFieldName
                                        << kUserDefinedTimeName
@@ -155,18 +154,17 @@ TEST_F(InternalUnpackBucketStageTest, UnpackEmptyInclude) {
         expCtx);
     unpack->setSource(source.get());
 
-    // We should produce metadata only, one per measurement in the bucket.
+    // We should produce empty documents, one per measurement in the bucket.
     for (auto idx = 0; idx < 2; ++idx) {
         auto next = unpack->getNext();
         ASSERT_TRUE(next.isAdvanced());
-        ASSERT_DOCUMENT_EQ(next.getDocument(), Document(fromjson("{myMeta: {m1: 999, m2: 9999}}")));
+        ASSERT_DOCUMENT_EQ(next.getDocument(), Document(fromjson("{}")));
     }
 
     for (auto idx = 0; idx < 2; ++idx) {
         auto next = unpack->getNext();
         ASSERT_TRUE(next.isAdvanced());
-        ASSERT_DOCUMENT_EQ(next.getDocument(),
-                           Document(fromjson("{myMeta: {m1: 9, m2: 9, m3: 9}}")));
+        ASSERT_DOCUMENT_EQ(next.getDocument(), Document(fromjson("{}")));
     }
 
     auto next = unpack->getNext();
@@ -254,8 +252,7 @@ TEST_F(InternalUnpackBucketStageTest, UnpackBasicIncludeWithDollarPrefix) {
 
     auto spec = BSON("$_internalUnpackBucket"
                      << BSON("include" << BSON_ARRAY("_id"
-                                                     << "time"
-                                                     << "$a"
+                                                     << "time" << kUserDefinedMetaName << "$a"
                                                      << "b")
                                        << DocumentSourceInternalUnpackBucket::kTimeFieldName
                                        << kUserDefinedTimeName
@@ -401,13 +398,60 @@ TEST_F(InternalUnpackBucketStageTest,
              << BSON("exclude" << BSONArray() << DocumentSourceInternalUnpackBucket::kTimeFieldName
                                << kUserDefinedTimeName));
     auto unpack = DocumentSourceInternalUnpackBucket::createFromBson(spec.firstElement(), expCtx);
-
     auto source =
         DocumentSourceMock::createForTest({"{data: {_id: {'1':1, "
                                            "'0':2, '2': 3}, time: {'1':1, '0': 2, '2': 3}}}",
                                            "{data: {_id: {'1':4, "
                                            "'0':5, '2':6}, time: {'1':4, '0': 5, '2': 6}}}"},
                                           expCtx);
+
+    unpack->setSource(source.get());
+
+    auto next = unpack->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    ASSERT_DOCUMENT_EQ(next.getDocument(), Document(fromjson("{time: 1, _id: 1}")));
+
+    next = unpack->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    ASSERT_DOCUMENT_EQ(next.getDocument(), Document(fromjson("{time: 2, _id: 2}")));
+
+    next = unpack->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    ASSERT_DOCUMENT_EQ(next.getDocument(), Document(fromjson("{time: 3, _id: 3}")));
+
+    // Second bucket
+    next = unpack->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    ASSERT_DOCUMENT_EQ(next.getDocument(), Document(fromjson("{time: 4, _id: 4}")));
+
+    next = unpack->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    ASSERT_DOCUMENT_EQ(next.getDocument(), Document(fromjson("{time: 5, _id: 5}")));
+
+    next = unpack->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    ASSERT_DOCUMENT_EQ(next.getDocument(), Document(fromjson("{time: 6, _id: 6}")));
+
+    next = unpack->getNext();
+    ASSERT_TRUE(next.isEOF());
+}
+
+TEST_F(InternalUnpackBucketStageTest, BucketUnpackerHandlesExcludedMetadataWhenBucketHasMetadata) {
+    auto expCtx = getExpCtx();
+    auto spec = BSON("$_internalUnpackBucket"
+                     << BSON("exclude" << BSON_ARRAY(kUserDefinedMetaName)
+                                       << DocumentSourceInternalUnpackBucket::kTimeFieldName
+                                       << kUserDefinedTimeName
+                                       << DocumentSourceInternalUnpackBucket::kMetaFieldName
+                                       << kUserDefinedMetaName));
+    auto unpack = DocumentSourceInternalUnpackBucket::createFromBson(spec.firstElement(), expCtx);
+    auto source = DocumentSourceMock::createForTest(
+        {"{meta: {'m1': 999, 'm2': 9999}, data: {_id: {'1':1, "
+         "'0':2, '2': 3}, time: {'1':1, '0': 2, '2': 3}}}",
+         "{meta: {'m1': 9, 'm2': 9, 'm3': 9}, data: {_id: {'1':4, "
+         "'0':5, '2':6}, time: {'1':4, '0': 5, '2': 6}}}"},
+        expCtx);
+
     unpack->setSource(source.get());
 
     auto next = unpack->getNext();
@@ -454,10 +498,10 @@ TEST_F(InternalUnpackBucketStageTest, BucketUnpackerThrowsOnUndefinedMetadata) {
                                            "'0':2, '2': 3}, time: {'1':1, '0': 2, '2': 3}}}"},
                                           expCtx);
     unpack->setSource(source.get());
-    ASSERT_THROWS_CODE(unpack->getNext(), AssertionException, 5346511);
+    ASSERT_THROWS_CODE(unpack->getNext(), AssertionException, 5369600);
 }
 
-TEST_F(InternalUnpackBucketStageTest, BucketUnpackerThrowsOnMissingMetadata) {
+TEST_F(InternalUnpackBucketStageTest, BucketUnpackerThrowsOnMissingMetadataWhenExpectedInBuckets) {
     auto expCtx = getExpCtx();
     auto spec =
         BSON("$_internalUnpackBucket"
@@ -472,9 +516,27 @@ TEST_F(InternalUnpackBucketStageTest, BucketUnpackerThrowsOnMissingMetadata) {
                                            "'0':2, '2': 3}, time: {'1':1, '0': 2, '2': 3}}}"},
                                           expCtx);
     unpack->setSource(source.get());
-    ASSERT_THROWS_CODE(unpack->getNext(), AssertionException, 5346511);
+    ASSERT_THROWS_CODE(unpack->getNext(), AssertionException, 5369600);
 }
 
+TEST_F(InternalUnpackBucketStageTest, BucketUnpackerThrowsWhenMetadataIsPresentUnexpectedly) {
+    auto expCtx = getExpCtx();
+    auto spec =
+        BSON("$_internalUnpackBucket"
+             << BSON("exclude" << BSONArray() << DocumentSourceInternalUnpackBucket::kTimeFieldName
+                               << kUserDefinedTimeName));
+    auto unpack = DocumentSourceInternalUnpackBucket::createFromBson(spec.firstElement(), expCtx);
+
+    auto source =
+        DocumentSourceMock::createForTest({"{meta: {'m1': 999, 'm2': 9999}, data: {_id: {'1':1, "
+                                           "'0':2, '2': 3}, time: {'1':1, '0': 2, '2': 3}}}",
+                                           "{meta: null, data: {_id: {'1':4, "
+                                           "'0':5, '2':6}, time: {'1':4, '0': 5, '2': 6}}}"},
+                                          expCtx);
+    unpack->setSource(source.get());
+
+    ASSERT_THROWS_CODE(unpack->getNext(), AssertionException, 5369601);
+}
 
 TEST_F(InternalUnpackBucketStageTest, BucketUnpackerHandlesNullMetadata) {
     auto expCtx = getExpCtx();