SERVER-32288 Add isAssignedShardKey field to CollectionType

3 files changed, 62 insertions, 1 deletions

diff --git a/src/mongo/s/catalog/type_collection.cpp b/src/mongo/s/catalog/type_collection.cpp
index f071f2f42c9..6c9a0644327 100644
--- a/src/mongo/s/catalog/type_collection.cpp
+++ b/src/mongo/s/catalog/type_collection.cpp
@@ -42,6 +42,7 @@ namespace {
 
 const BSONField<bool> kNoBalance("noBalance");
 const BSONField<bool> kDropped("dropped");
+const auto kIsAssignedShardKey = "isAssignedShardKey"_sd;
 
 }  // namespace
 
@@ -178,6 +179,18 @@ StatusWith<CollectionType> CollectionType::fromBSON(const BSONObj& source) {
         }
     }
 
+    {
+        bool isAssignedShardKey;
+        Status status = bsonExtractBooleanField(source, kIsAssignedShardKey, &isAssignedShardKey);
+        if (status.isOK()) {
+            coll._isAssignedShardKey = isAssignedShardKey;
+        } else if (status == ErrorCodes::NoSuchKey) {
+            // isAssignedShardKey can be missing in which case it is presumed as true.
+        } else {
+            return status;
+        }
+    }
+
     return StatusWith<CollectionType>(coll);
 }
 
@@ -251,6 +264,10 @@ BSONObj CollectionType::toBSON() const {
         builder.append(kNoBalance.name(), !_allowBalance.get());
     }
 
+    if (_isAssignedShardKey) {
+        builder.append(kIsAssignedShardKey, !_isAssignedShardKey.get());
+    }
+
     return builder.obj();
 }
 
diff --git a/src/mongo/s/catalog/type_collection.h b/src/mongo/s/catalog/type_collection.h
index af6247b5232..796b4f7707a 100644
--- a/src/mongo/s/catalog/type_collection.h
+++ b/src/mongo/s/catalog/type_collection.h
@@ -62,7 +62,8 @@ class StatusWith;
  *      },
  *      "unique" : false,
  *      "uuid" : UUID,
- *      "noBalance" : false
+ *      "noBalance" : false,
+ *      "allowSplit" : false
  *   }
  *
  */
@@ -157,6 +158,10 @@ public:
         return _allowBalance.get_value_or(true);
     }
 
+    bool isAssignedShardKey() const {
+        return _isAssignedShardKey.get_value_or(true);
+    }
+
     bool hasSameOptions(CollectionType& other);
 
 private:
@@ -186,6 +191,10 @@ private:
 
     // Optional whether balancing is allowed for this collection. If missing, implies true.
     boost::optional<bool> _allowBalance;
+
+    // Optional whether user has assigned a shard key to this collection before.
+    // Implicitly true if missing.
+    boost::optional<bool> _isAssignedShardKey;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/s/catalog/type_collection_test.cpp b/src/mongo/s/catalog/type_collection_test.cpp
index 89ad96733b9..39b2ea84929 100644
--- a/src/mongo/s/catalog/type_collection_test.cpp
+++ b/src/mongo/s/catalog/type_collection_test.cpp
@@ -69,6 +69,41 @@ TEST(CollectionType, Basic) {
     ASSERT_EQUALS(coll.getUnique(), true);
     ASSERT_EQUALS(coll.getAllowBalance(), true);
     ASSERT_EQUALS(coll.getDropped(), false);
+    ASSERT_EQUALS(coll.isAssignedShardKey(), true);
+}
+
+TEST(CollectionType, AllFieldsPresent) {
+    const OID oid = OID::gen();
+    const auto uuid = UUID::gen();
+    StatusWith<CollectionType> status =
+        CollectionType::fromBSON(BSON(CollectionType::fullNs("db.coll")
+                                      << CollectionType::epoch(oid)
+                                      << CollectionType::updatedAt(Date_t::fromMillisSinceEpoch(1))
+                                      << CollectionType::keyPattern(BSON("a" << 1))
+                                      << CollectionType::defaultCollation(BSON("locale"
+                                                                               << "fr_CA"))
+                                      << CollectionType::unique(true)
+                                      << CollectionType::uuid()
+                                      << uuid
+                                      << "isAssignedShardKey"
+                                      << false));
+    ASSERT_TRUE(status.isOK());
+
+    CollectionType coll = status.getValue();
+    ASSERT_TRUE(coll.validate().isOK());
+    ASSERT(coll.getNs() == NamespaceString{"db.coll"});
+    ASSERT_EQUALS(coll.getEpoch(), oid);
+    ASSERT_EQUALS(coll.getUpdatedAt(), Date_t::fromMillisSinceEpoch(1));
+    ASSERT_BSONOBJ_EQ(coll.getKeyPattern().toBSON(), BSON("a" << 1));
+    ASSERT_BSONOBJ_EQ(coll.getDefaultCollation(),
+                      BSON("locale"
+                           << "fr_CA"));
+    ASSERT_EQUALS(coll.getUnique(), true);
+    ASSERT_EQUALS(coll.getAllowBalance(), true);
+    ASSERT_EQUALS(coll.getDropped(), false);
+    ASSERT_TRUE(coll.getUUID());
+    ASSERT_EQUALS(*coll.getUUID(), uuid);
+    ASSERT_EQUALS(coll.isAssignedShardKey(), false);
 }
 
 TEST(CollectionType, EmptyDefaultCollationFailsToParse) {