SERVER-56227 Add user-facing command for permitMigrations using DDL coordinator

(cherry picked from commit ddb38e8c0da429630eebe340f6bbd487d2730242)

15 files changed, 377 insertions, 35 deletions

diff --git a/jstests/core/views/views_all_commands.js b/jstests/core/views/views_all_commands.js
index 9f3d658c241..f8d1be30b7b 100644
--- a/jstests/core/views/views_all_commands.js
+++ b/jstests/core/views/views_all_commands.js
@@ -108,6 +108,7 @@ let viewsCommandTests = {
     _recvChunkCommit: {skip: isAnInternalCommand},
     _recvChunkStart: {skip: isAnInternalCommand},
     _recvChunkStatus: {skip: isAnInternalCommand},
+    _shardsvrSetAllowMigrations: {skip: isAnInternalCommand},
     _shardsvrShardCollection: {skip: isAnInternalCommand},
     _transferMods: {skip: isAnInternalCommand},
     abortTransaction: {skip: isUnrelated},
@@ -463,6 +464,16 @@ let viewsCommandTests = {
     saslContinue: {skip: isUnrelated},
     saslStart: {skip: isUnrelated},
     serverStatus: {command: {serverStatus: 1}, skip: isUnrelated},
+    setAllowMigrations: {
+        command: {setAllowMigrations: "test.view", allowMigrations: false},
+        setup: function(conn) {
+            assert.commandWorked(conn.adminCommand({enableSharding: "test"}));
+        },
+        expectedErrorCode: ErrorCodes.NamespaceNotSharded,
+        skipStandalone: true,
+        expectFailure: true,
+        isAdminCommand: true
+    },
     setCommittedSnapshot: {skip: isAnInternalCommand},
     setFeatureCompatibilityVersion: {skip: isUnrelated},
     setFreeMonitoring: {skip: isUnrelated},
diff --git a/jstests/replsets/db_reads_while_recovering_all_commands.js b/jstests/replsets/db_reads_while_recovering_all_commands.js
index a3c808499d5..05dade63ca3 100644
--- a/jstests/replsets/db_reads_while_recovering_all_commands.js
+++ b/jstests/replsets/db_reads_while_recovering_all_commands.js
@@ -67,6 +67,7 @@ const allCommands = {
     _recvChunkStatus: {skip: isPrimaryOnly},
     _shardsvrCloneCatalogData: {skip: isPrimaryOnly},
     _shardsvrMovePrimary: {skip: isPrimaryOnly},
+    _shardsvrSetAllowMigrations: {skip: isPrimaryOnly},
     _shardsvrShardCollection: {skip: isPrimaryOnly},
     _transferMods: {skip: isPrimaryOnly},
     abortTransaction: {skip: isPrimaryOnly},
diff --git a/jstests/sharding/database_and_shard_versioning_all_commands.js b/jstests/sharding/database_and_shard_versioning_all_commands.js
index 4072dc910a5..52a5bd35fe7 100644
--- a/jstests/sharding/database_and_shard_versioning_all_commands.js
+++ b/jstests/sharding/database_and_shard_versioning_all_commands.js
@@ -386,6 +386,7 @@ let testCases = {
     saslContinue: {skip: "not on a user database"},
     saslStart: {skip: "not on a user database"},
     serverStatus: {skip: "executes locally on mongos (not sent to any remote node)"},
+    setAllowMigrations: {skip: "not on a user database"},
     setFeatureCompatibilityVersion: {skip: "not on a user database"},
     setFreeMonitoring:
         {skip: "explicitly fails for mongos, primary mongod only", conditional: true},
diff --git a/jstests/sharding/move_chunk_permitMigrations.js b/jstests/sharding/move_chunk_permitMigrations.js
index 3966ea0229b..dc77b3144ca 100644
--- a/jstests/sharding/move_chunk_permitMigrations.js
+++ b/jstests/sharding/move_chunk_permitMigrations.js
@@ -3,7 +3,8 @@
  * moveChunk and disables the balancer.
  *
  * @tags: [
- *   multiversion_incompatible
+ *   multiversion_incompatible,
+ *   does_not_support_stepdowns,
  * ]
  */
 (function() {
@@ -14,7 +15,7 @@ load('jstests/libs/parallel_shell_helpers.js');
 
 const st = new ShardingTest({shards: 2});
 const configDB = st.s.getDB("config");
-const dbName = 'PermitMigrations';
+const dbName = 'AllowMigrations';
 
 // Resets database dbName and enables sharding and establishes shard0 as primary, test case agnostic
 const setUpDb = function setUpDatabaseAndEnableSharding() {
@@ -23,15 +24,13 @@ const setUpDb = function setUpDatabaseAndEnableSharding() {
         st.s.adminCommand({enableSharding: dbName, primaryShard: st.shard0.shardName}));
 };
 
-const setPermitMigrations = function(ns, permit) {
-    // For now update the flag manually, a user-facing command will be implemented with
-    // SERVER-56227.
-    assert.commandWorked(configDB.collections.updateOne(
-        {_id: ns}, {$set: {permitMigrations: permit}}, {writeConcern: {w: "majority"}}));
+// Use the setAllowMigrations command to set the permitMigrations flag in the collection.
+const setAllowMigrationsCmd = function(ns, allow) {
+    assert.commandWorked(st.s.adminCommand({setAllowMigrations: ns, allowMigrations: allow}));
 };
 
-// Tests that moveChunk does not succeed when {permitMigrations: false}
-(function testPermitMigrationsFalsePreventsMoveChunk() {
+// Tests that moveChunk does not succeed when setAllowMigrations is called with a false value.
+(function testSetAllowMigrationsFalsePreventsMoveChunk() {
     setUpDb();
 
     const collName = "collA";
@@ -41,31 +40,19 @@ const setPermitMigrations = function(ns, permit) {
     assert.commandWorked(st.s.getDB(dbName).getCollection(collName).insert({_id: 1}));
     assert.commandWorked(st.s.adminCommand({shardCollection: ns, key: {_id: 1}}));
 
-    // Confirm that an inProgress moveChunk fails once {permitMigrations: false}
-    const fp = configureFailPoint(st.shard0, "moveChunkHangAtStep4");
-    const awaitResult = startParallelShell(
-        funWithArgs(function(ns, toShardName) {
-            assert.commandFailedWithCode(
-                db.adminCommand({moveChunk: ns, find: {_id: 0}, to: toShardName}),
-                ErrorCodes.ConflictingOperationInProgress);
-        }, ns, st.shard1.shardName), st.s.port);
-    fp.wait();
-    setPermitMigrations(ns, false);
-    fp.off();
-    awaitResult();
-
-    // {permitMigrations: false} is set, sending a new moveChunk command should also fail.
+    setAllowMigrationsCmd(ns, false);
+
+    // setAllowMigrations was called, sending a new moveChunk command should fail.
     assert.commandFailedWithCode(
         st.s.adminCommand({moveChunk: ns, find: {_id: 0}, to: st.shard1.shardName}),
         ErrorCodes.ConflictingOperationInProgress);
 })();
 
-// Tests {permitMigrations: false} disables balancing for collB and does not interfere with
+// Tests setAllowMigrations disables balancing for collB and does not interfere with
 // balancing for collA.
 //
 // collBSetParams specify the field(s) that will be set on the collB in config.collections.
-const testBalancer = function testPermitMigrationsFalseDisablesBalancer(permitMigrations,
-                                                                        collBSetNoBalanceParam) {
+const testBalancer = function(setAllowMigrations, collBSetNoBalanceParam) {
     setUpDb();
 
     const collAName = "collA";
@@ -92,11 +79,12 @@ const testBalancer = function testPermitMigrationsFalseDisablesBalancer(permitMi
             configDB.chunks.countDocuments({ns: coll.getFullName(), shard: st.shard0.shardName}));
     }
 
-    jsTestLog(`Disabling balancing of ${collB.getFullName()} with permitMigrations ${
-        permitMigrations} and parameters ${tojson(collBSetNoBalanceParam)}`);
+    jsTestLog(`Disabling balancing of ${collB.getFullName()} with setAllowMigrations ${
+        setAllowMigrations} and parameters ${tojson(collBSetNoBalanceParam)}`);
     assert.commandWorked(
         configDB.collections.update({_id: collB.getFullName()}, {$set: collBSetNoBalanceParam}));
-    setPermitMigrations(collB.getFullName(), permitMigrations);
+
+    setAllowMigrationsCmd(collB.getFullName(), setAllowMigrations);
 
     st.startBalancer();
     assert.soon(() => {
@@ -110,14 +98,40 @@ const testBalancer = function testPermitMigrationsFalseDisablesBalancer(permitMi
     }, `Balancer failed to balance ${collA.getFullName()}`, 1000 * 60 * 10);
     st.stopBalancer();
 
+    // collB should still be unbalanced.
     assert.eq(
         4, configDB.chunks.countDocuments({ns: collB.getFullName(), shard: st.shard0.shardName}));
 };
 
+const testSetAllowMigrationsCommand = function() {
+    setUpDb();
+
+    const collName = "foo";
+    const ns = dbName + "." + collName;
+
+    assert.commandWorked(st.s.adminCommand({shardCollection: ns, key: {x: 1}}));
+
+    // Use setAllowMigrations to forbid migrations from happening
+    setAllowMigrationsCmd(ns, false);
+
+    // Check that allowMigrations has been set to 'false' on the configsvr config.collections.
+    assert.eq(false, configDB.collections.findOne({_id: ns}).permitMigrations);
+
+    // Use setAllowMigrations to allow migrations to happen
+    setAllowMigrationsCmd(ns, true);
+
+    // Check that permitMigrations has been unset (that implies migrations are allowed) on the
+    // configsvr config.collections.
+    assert.eq(undefined, configDB.collections.findOne({_id: ns}).permitMigrations);
+};
+
 // Test cases that should disable the balancer.
-testBalancer(false /* permitMigrations */, {});
-testBalancer(false /* permitMigrations */, {noBalance: false});
-testBalancer(false /* permitMigrations */, {noBalance: true});
+testBalancer(false /* setAllowMigrations */, {});
+testBalancer(false /* setAllowMigrations */, {noBalance: false});
+testBalancer(false /* setAllowMigrations */, {noBalance: true});
+
+// Test the setAllowMigrations command.
+testSetAllowMigrationsCommand();
 
 st.stop();
 })();
diff --git a/jstests/sharding/safe_secondary_reads_drop_recreate.js b/jstests/sharding/safe_secondary_reads_drop_recreate.js
index 6e826e8e565..575746f83c1 100644
--- a/jstests/sharding/safe_secondary_reads_drop_recreate.js
+++ b/jstests/sharding/safe_secondary_reads_drop_recreate.js
@@ -279,6 +279,7 @@ let testCases = {
     saslContinue: {skip: "primary only"},
     saslStart: {skip: "primary only"},
     serverStatus: {skip: "does not return user data"},
+    setAllowMigrations: {skip: "primary only"},
     setCommittedSnapshot: {skip: "does not return user data"},
     setFeatureCompatibilityVersion: {skip: "primary only"},
     setFreeMonitoring: {skip: "primary only"},
diff --git a/jstests/sharding/safe_secondary_reads_single_migration_suspend_range_deletion.js b/jstests/sharding/safe_secondary_reads_single_migration_suspend_range_deletion.js
index d46e09e1a3f..88318a11d31 100644
--- a/jstests/sharding/safe_secondary_reads_single_migration_suspend_range_deletion.js
+++ b/jstests/sharding/safe_secondary_reads_single_migration_suspend_range_deletion.js
@@ -314,6 +314,7 @@ let testCases = {
     saslContinue: {skip: "primary only"},
     saslStart: {skip: "primary only"},
     serverStatus: {skip: "does not return user data"},
+    setAllowMigrations: {skip: "primary only"},
     setCommittedSnapshot: {skip: "does not return user data"},
     setFeatureCompatibilityVersion: {skip: "primary only"},
     setFreeMonitoring: {skip: "primary only"},
diff --git a/jstests/sharding/safe_secondary_reads_single_migration_waitForDelete.js b/jstests/sharding/safe_secondary_reads_single_migration_waitForDelete.js
index 2ca7bc352d0..9469bb43581 100644
--- a/jstests/sharding/safe_secondary_reads_single_migration_waitForDelete.js
+++ b/jstests/sharding/safe_secondary_reads_single_migration_waitForDelete.js
@@ -284,6 +284,7 @@ let testCases = {
     saslContinue: {skip: "primary only"},
     saslStart: {skip: "primary only"},
     serverStatus: {skip: "does not return user data"},
+    setAllowMigrations: {skip: "primary only"},
     setCommittedSnapshot: {skip: "does not return user data"},
     setFeatureCompatibilityVersion: {skip: "primary only"},
     setFreeMonitoring: {skip: "primary only"},
diff --git a/src/mongo/db/s/SConscript b/src/mongo/db/s/SConscript
index 37e14c2ffce..e4418131d92 100644
--- a/src/mongo/db/s/SConscript
+++ b/src/mongo/db/s/SConscript
@@ -346,6 +346,7 @@ env.Library(
         'set_shard_version_command.cpp',
         'sharding_server_status.cpp',
         'sharding_state_command.cpp',
+        'shardsvr_set_allow_migrations_command.cpp',
         'shardsvr_shard_collection.cpp',
         'split_chunk_command.cpp',
         'split_vector_command.cpp',
diff --git a/src/mongo/db/s/shardsvr_set_allow_migrations_command.cpp b/src/mongo/db/s/shardsvr_set_allow_migrations_command.cpp
new file mode 100644
index 00000000000..244653e2783
--- /dev/null
+++ b/src/mongo/db/s/shardsvr_set_allow_migrations_command.cpp
@@ -0,0 +1,142 @@
+/**
+ *    Copyright (C) 2021-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kSharding
+
+#include "mongo/db/auth/authorization_session.h"
+#include "mongo/db/commands.h"
+#include "mongo/db/s/sharding_state.h"
+#include "mongo/s/grid.h"
+#include "mongo/s/request_types/set_allow_migrations_gen.h"
+#include "mongo/s/write_ops/batched_command_request.h"
+#include "mongo/s/write_ops/batched_command_response.h"
+
+namespace mongo {
+namespace {
+
+class ShardsvrSetAllowMigrationsCommand final
+    : public TypedCommand<ShardsvrSetAllowMigrationsCommand> {
+public:
+    using Request = ShardsvrSetAllowMigrations;
+
+    std::string help() const override {
+        return "Internal command. Do not call directly. Enable/disable migrations in a collection.";
+    }
+
+    bool adminOnly() const override {
+        return false;
+    }
+
+    AllowedOnSecondary secondaryAllowed(ServiceContext*) const override {
+        return AllowedOnSecondary::kNever;
+    }
+
+    class Invocation final : public InvocationBase {
+    public:
+        using InvocationBase::InvocationBase;
+
+        bool isCollectionSharded(OperationContext* opCtx, const NamespaceString& nss) {
+            try {
+                uassertStatusOK(Grid::get(opCtx)->catalogClient()->getCollection(opCtx, nss));
+                return true;
+            } catch (ExceptionFor<ErrorCodes::NamespaceNotFound>&) {
+                // The collection is unsharded or doesn't exist
+                return false;
+            }
+        }
+
+        void typedRun(OperationContext* opCtx) {
+            uassertStatusOK(ShardingState::get(opCtx)->canAcceptShardedCommands());
+
+            opCtx->setAlwaysInterruptAtStepDownOrUp();
+
+            uassert(ErrorCodes::InvalidOptions,
+                    str::stream() << Request::kCommandName
+                                  << " must be called with majority writeConcern, got "
+                                  << request().toBSON(BSONObj()),
+                    opCtx->getWriteConcern().wMode == WriteConcernOptions::kMajority);
+
+            const auto& nss = ns();
+
+            uassert(ErrorCodes::NamespaceNotSharded,
+                    "Collection must be sharded so migrations can be blocked",
+                    isCollectionSharded(opCtx, nss));
+
+            const auto configShard = Grid::get(opCtx)->shardRegistry()->getConfigShard();
+
+            BatchedCommandRequest updateRequest([&]() {
+                write_ops::Update updateOp(CollectionType::ConfigNS);
+                updateOp.setUpdates({[&] {
+                    write_ops::UpdateOpEntry entry;
+                    entry.setQ(BSON(CollectionType::fullNs << nss.ns()));
+                    if (request().getAllowMigrations()) {
+                        entry.setU(write_ops::UpdateModification(BSON(
+                            "$unset" << BSON(CollectionType::permitMigrations.name() << true))));
+                    } else {
+                        entry.setU(write_ops::UpdateModification(BSON(
+                            "$set" << BSON(CollectionType::permitMigrations.name() << false))));
+                    }
+                    entry.setMulti(false);
+                    return entry;
+                }()});
+                return updateOp;
+            }());
+
+            updateRequest.setWriteConcern(ShardingCatalogClient::kMajorityWriteConcern.toBSON());
+
+            auto response = configShard->runBatchWriteCommand(opCtx,
+                                                              Shard::kDefaultConfigCommandTimeout,
+                                                              updateRequest,
+                                                              Shard::RetryPolicy::kIdempotent);
+
+            uassertStatusOK(response.toStatus());
+        }
+
+    private:
+        NamespaceString ns() const override {
+            return request().getNamespace();
+        }
+
+        bool supportsWriteConcern() const override {
+            return true;
+        }
+
+        void doCheckAuthorization(OperationContext* opCtx) const override {
+            uassert(ErrorCodes::Unauthorized,
+                    "Unauthorized",
+                    AuthorizationSession::get(opCtx->getClient())
+                        ->isAuthorizedForActionsOnResource(ResourcePattern::forClusterResource(),
+                                                           ActionType::internal));
+        }
+    };
+
+} shardsvrSetAllowMigrationsCommand;
+
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/s/SConscript b/src/mongo/s/SConscript
index 807986aca0c..fe4c77bea88 100644
--- a/src/mongo/s/SConscript
+++ b/src/mongo/s/SConscript
@@ -190,6 +190,7 @@ env.Library(
         env.Idlc('request_types/move_primary.idl')[0],
         env.Idlc('request_types/shard_collection.idl')[0],
         env.Idlc('request_types/clone_collection_options_from_primary_shard.idl')[0],
+        env.Idlc('request_types/set_allow_migrations.idl')[0],
         env.Idlc('request_types/wait_for_fail_point.idl')[0],
     ],
     LIBDEPS=[
diff --git a/src/mongo/s/catalog/type_collection.cpp b/src/mongo/s/catalog/type_collection.cpp
index 884aef938f2..1cd2e2e0de7 100644
--- a/src/mongo/s/catalog/type_collection.cpp
+++ b/src/mongo/s/catalog/type_collection.cpp
@@ -44,7 +44,6 @@ namespace {
 const BSONField<bool> kNoBalance("noBalance");
 const BSONField<bool> kDropped("dropped");
 const auto kIsAssignedShardKey = "isAssignedShardKey"_sd;
-const BSONField<bool> kPermitMigrations("permitMigrations");
 
 }  // namespace
 
@@ -57,6 +56,7 @@ const BSONField<BSONObj> CollectionType::keyPattern("key");
 const BSONField<BSONObj> CollectionType::defaultCollation("defaultCollation");
 const BSONField<bool> CollectionType::unique("unique");
 const BSONField<UUID> CollectionType::uuid("uuid");
+const BSONField<bool> CollectionType::permitMigrations("permitMigrations");
 
 StatusWith<CollectionType> CollectionType::fromBSON(const BSONObj& source) {
     CollectionType coll;
@@ -196,7 +196,7 @@ StatusWith<CollectionType> CollectionType::fromBSON(const BSONObj& source) {
     {
         bool collPermitMigrations;
         Status status =
-            bsonExtractBooleanField(source, kPermitMigrations.name(), &collPermitMigrations);
+            bsonExtractBooleanField(source, permitMigrations.name(), &collPermitMigrations);
         if (status.isOK()) {
             coll._permitMigrations = collPermitMigrations;
         } else if (status == ErrorCodes::NoSuchKey) {
@@ -284,7 +284,7 @@ BSONObj CollectionType::toBSON() const {
     }
 
     if (_permitMigrations.is_initialized()) {
-        builder.append(kPermitMigrations.name(), _permitMigrations.get());
+        builder.append(permitMigrations.name(), _permitMigrations.get());
     }
 
     return builder.obj();
diff --git a/src/mongo/s/catalog/type_collection.h b/src/mongo/s/catalog/type_collection.h
index 0e138402c51..89d19d4c2f9 100644
--- a/src/mongo/s/catalog/type_collection.h
+++ b/src/mongo/s/catalog/type_collection.h
@@ -81,6 +81,7 @@ public:
     static const BSONField<BSONObj> defaultCollation;
     static const BSONField<bool> unique;
     static const BSONField<UUID> uuid;
+    static const BSONField<bool> permitMigrations;
 
     /**
      * Constructs a new DatabaseType object from BSON. Also does validation of the contents.
diff --git a/src/mongo/s/commands/SConscript b/src/mongo/s/commands/SConscript
index ed74f00aabd..f363155acee 100644
--- a/src/mongo/s/commands/SConscript
+++ b/src/mongo/s/commands/SConscript
@@ -75,6 +75,7 @@ env.Library(
         'cluster_repl_set_get_status_cmd.cpp',
         'cluster_reset_error_cmd.cpp',
         'cluster_restart_catalog_command.cpp',
+        'cluster_set_allow_migrations_cmd.cpp',
         'cluster_set_feature_compatibility_version_cmd.cpp',
         'cluster_set_free_monitoring.cpp' if get_option("enable-free-mon") == 'on' else [],
         'cluster_shard_collection_cmd.cpp',
diff --git a/src/mongo/s/commands/cluster_set_allow_migrations_cmd.cpp b/src/mongo/s/commands/cluster_set_allow_migrations_cmd.cpp
new file mode 100644
index 00000000000..bffdc7e1014
--- /dev/null
+++ b/src/mongo/s/commands/cluster_set_allow_migrations_cmd.cpp
@@ -0,0 +1,104 @@
+/**
+ *    Copyright (C) 2018-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kSharding
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/rpc/get_status_from_command_result.h"
+
+#include "mongo/db/auth/authorization_session.h"
+#include "mongo/db/commands.h"
+#include "mongo/s/cluster_commands_helpers.h"
+#include "mongo/s/grid.h"
+#include "mongo/s/request_types/set_allow_migrations_gen.h"
+
+namespace mongo {
+namespace {
+
+class SetAllowMigrationsCmd final : public TypedCommand<SetAllowMigrationsCmd> {
+public:
+    using Request = SetAllowMigrations;
+
+    AllowedOnSecondary secondaryAllowed(ServiceContext*) const override {
+        return AllowedOnSecondary::kNever;
+    }
+
+    bool adminOnly() const override {
+        return true;
+    }
+
+    class Invocation final : public InvocationBase {
+    public:
+        using InvocationBase::InvocationBase;
+
+        void typedRun(OperationContext* opCtx) {
+            const auto& nss = ns();
+
+            SetAllowMigrationsRequest allowMigrationsRequest;
+            allowMigrationsRequest.setAllowMigrations(request().getAllowMigrations());
+            ShardsvrSetAllowMigrations shardsvrRequest(nss);
+            shardsvrRequest.setSetAllowMigrationsRequest(allowMigrationsRequest);
+
+            auto catalogCache = Grid::get(opCtx)->catalogCache();
+            const auto dbInfo = uassertStatusOK(catalogCache->getDatabase(opCtx, nss.db()));
+            auto cmdResponse = executeCommandAgainstDatabasePrimary(
+                opCtx,
+                nss.db(),
+                dbInfo,
+                CommandHelpers::appendMajorityWriteConcern(shardsvrRequest.toBSON({})),
+                ReadPreferenceSetting(ReadPreference::PrimaryOnly),
+                Shard::RetryPolicy::kIdempotent);
+
+            const auto remoteResponse = uassertStatusOK(cmdResponse.swResponse);
+            uassertStatusOK(getStatusFromCommandResult(remoteResponse.data));
+        }
+
+        NamespaceString ns() const override {
+            return request().getCommandParameter();
+        }
+
+        // Considering this command will stop migrations, it is reasonable to ensure the same
+        // permissions as moveChunk.
+        void doCheckAuthorization(OperationContext* opCtx) const override {
+            uassert(ErrorCodes::Unauthorized,
+                    "Unauthorized to perform migration operations",
+                    AuthorizationSession::get(opCtx->getClient())
+                        ->isAuthorizedForActionsOnResource(ResourcePattern::forExactNamespace(ns()),
+                                                           ActionType::moveChunk));
+        }
+
+        bool supportsWriteConcern() const override {
+            return true;
+        }
+    };
+} setAllowMigrationsCmd;
+
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/s/request_types/set_allow_migrations.idl b/src/mongo/s/request_types/set_allow_migrations.idl
new file mode 100644
index 00000000000..9d1b74c4cf6
--- /dev/null
+++ b/src/mongo/s/request_types/set_allow_migrations.idl
@@ -0,0 +1,62 @@
+# Copyright(C) 2021 - present MongoDB, Inc.
+#
+# This program is free software : you can redistribute it and / or modify
+# it under the terms of the Server Side Public License, version 1,
+# as published by MongoDB, Inc.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.See the
+# Server Side Public License for more details.
+#
+# You should have received a copy of the Server Side Public License
+# along with this program.If not, see
+# < http://www.mongodb.com/licensing/server-side-public-license>.
+#
+# As a special exception, the copyright holders give permission to link the
+# code of portions of this program with the OpenSSL library under certain
+# conditions as described in each individual source file and distribute
+# linked combinations including the program with the OpenSSL library.You
+# must comply with the Server Side Public License in all respects for
+# all of the code used other than as permitted herein.If you modify file(s)
+# with this exception, you may extend this exception to your version of the
+# file(s), but you are not obligated to do so.If you do not wish to do so,
+# delete this exception statement from your version.If you delete this
+# exception statement from all source files in the program, then also delete
+# it in the license file.
+#
+
+global:
+    cpp_namespace: "mongo"
+
+imports:
+    - "mongo/idl/basic_types.idl"
+
+structs:
+    SetAllowMigrationsRequest:
+        description: "Parameters sent for the set allow migrations command."
+        strict: false
+        fields:
+            allowMigrations:
+                type: bool
+                description: "If false balancer rounds should be disabled and migrations commit prohibited."
+                optional: false
+
+commands:
+    setAllowMigrations:
+        cpp_name: SetAllowMigrations
+        description: "User faced setAllowMigrations command."
+        namespace: type
+        strict: true
+        type: namespacestring
+        chained_structs:
+            SetAllowMigrationsRequest: SetAllowMigrationsRequest
+
+    _shardsvrSetAllowMigrations:
+        cpp_name: shardsvrSetAllowMigrations
+        description: "Internal setAllowMigrations command for a shard."
+        strict: false
+        namespace: concatenate_with_db
+        chained_structs:
+            SetAllowMigrationsRequest: SetAllowMigrationsRequest
+