SERVER-68513 The _configsvrRunRestore command should restore databases with unsharded collections

2 files changed, 87 insertions, 5 deletions

diff --git a/jstests/sharding/run_restore_unsharded.js b/jstests/sharding/run_restore_unsharded.js
new file mode 100644
index 00000000000..038e79f1dd9
--- /dev/null
+++ b/jstests/sharding/run_restore_unsharded.js
@@ -0,0 +1,67 @@
+/**
+ * Tests that the "_configsvrRunRestore" command restores databases with unsharded collections
+ * referenced in the "local.system.collections_to_restore" collection.
+ *
+ * @tags: [
+ *      requires_persistence,
+ * ]
+ */
+(function() {
+"use strict";
+
+load("jstests/libs/feature_flag_util.js");
+
+const s =
+    new ShardingTest({name: "runRestore", shards: 2, mongos: 1, config: 1, other: {chunkSize: 1}});
+
+let mongos = s.s0;
+let db = s.getDB("test");
+if (!FeatureFlagUtil.isEnabled(s.configRS.getPrimary().getDB("test"), "SelectiveBackup")) {
+    jsTestLog("Skipping as featureFlagSelectiveBackup is not enabled");
+    s.stop();
+    return;
+}
+
+s.adminCommand({enablesharding: "test"});
+s.ensurePrimaryShard("test", s.shard0.shardName);
+
+// Create an unsharded collection.
+assert.commandWorked(db.createCollection("a"));
+const collUUID =
+    s.shard0.getDB("test").runCommand({listCollections: 1}).cursor.firstBatch[0].info.uuid;
+
+// Only sharded collections appear in config.collections
+assert.eq(0, mongos.getDB("config").getCollection("collections").find({_id: "test.a"}).count());
+
+assert.eq(1, mongos.getDB("config").getCollection("locks").find({_id: "test"}).count());
+assert.eq(1, mongos.getDB("config").getCollection("databases").find({_id: "test"}).count());
+
+s.stop({noCleanData: true});
+
+const configDbPath = s.c0.dbpath;
+
+// Start the config server in standalone restore mode.
+let conn = MongoRunner.runMongod({noCleanData: true, dbpath: configDbPath, restore: ""});
+assert(conn);
+
+assert.commandWorked(conn.getDB("admin").runCommand({setParameter: 1, logLevel: 1}));
+
+// Create the "local.system.collections_to_restore" collection and insert "test.a".
+assert.commandWorked(conn.getDB("local").createCollection("system.collections_to_restore"));
+assert.commandWorked(conn.getDB("local").getCollection("system.collections_to_restore").insert({
+    ns: "test.a",
+    uuid: collUUID
+}));
+
+assert.commandWorked(conn.getDB("admin").runCommand({_configsvrRunRestore: 1}));
+
+// Only sharded collections appear in config.collections
+assert.eq(0, conn.getDB("config").getCollection("collections").find({_id: "test.a"}).count());
+
+let locks = conn.getDB("config").getCollection("locks").find({_id: "test"}).toArray();
+assert.eq(1, locks.length);
+assert.eq(0, locks[0].state);  // State::UNLOCKED
+assert.eq(1, conn.getDB("config").getCollection("databases").find({_id: "test"}).count());
+
+MongoRunner.stopMongod(conn);
+}());
diff --git a/src/mongo/db/s/config/configsvr_run_restore_command.cpp b/src/mongo/db/s/config/configsvr_run_restore_command.cpp
index 166c10a4e70..22e3f7b2985 100644
--- a/src/mongo/db/s/config/configsvr_run_restore_command.cpp
+++ b/src/mongo/db/s/config/configsvr_run_restore_command.cpp
@@ -73,6 +73,25 @@ ShouldRestoreDocument shouldRestoreDocument(OperationContext* opCtx,
                                                    : ShouldRestoreDocument::kNo;
 }
 
+std::set<std::string> getDatabasesToRestore(OperationContext* opCtx) {
+    auto findRequest = FindCommandRequest(NamespaceString::kConfigsvrRestoreNamespace);
+
+    std::set<std::string> databasesToRestore;
+    DBDirectClient client(opCtx);
+    auto it = client.find(findRequest);
+    while (it->more()) {
+        const auto doc = it->next();
+        if (!doc.hasField("ns")) {
+            continue;
+        }
+
+        NamespaceString nss(doc.getStringField("ns"));
+        databasesToRestore.emplace(nss.db());
+    }
+
+    return databasesToRestore;
+}
+
 // Modifications to this map should add new testing in 'sharded_backup_restore.js'.
 // { config collection namespace -> ( optional nss field name, optional UUID field name ) }
 const stdx::unordered_map<NamespaceString,
@@ -146,7 +165,7 @@ public:
 
         // Keeps track of database names for collections restored. Databases with no collections
         // restored will have their entries removed in the config collections.
-        std::set<std::string> databasesRestored;
+        std::set<std::string> databasesRestored = getDatabasesToRestore(opCtx);
 
         for (const auto& collectionEntry : kCollectionEntries) {
             const NamespaceString& nss = collectionEntry.first;
@@ -200,10 +219,6 @@ public:
                             "doc"_attr = doc,
                             "shouldRestore"_attr = shouldRestore);
 
-                if (shouldRestore == ShouldRestoreDocument::kYes && docNss) {
-                    databasesRestored.insert(docNss->db().toString());
-                }
-
                 if (shouldRestore == ShouldRestoreDocument::kYes ||
                     shouldRestore == ShouldRestoreDocument::kMaybe) {
                     continue;