feat: add sharded migration helpers MONGOSH-2139

packages/i18n/src/locales/en_US.ts

@@ -2124,6 +2124,18 @@ const translations: Catalog = {
                 'Disable balancing on a single collection in a sharded database. Does not affect balancing of other collections in a sharded cluster.',
               example: 'sh.disableBalancing(ns)',
             },
+            enableMigrations: {
+              link: 'https://mongodb.com/docs/manual/reference/method/sh.enableMigrations',
+              description:
+                'Enables migrations for a specific collection. Uses `setAllowMigrations` admin command.',
+              example: 'sh.enableMigrations(ns)',
+            },
+            disableMigrations: {
+              link: 'https://mongodb.com/docs/manual/reference/method/sh.disableMigrations',
+              description:
+                'Disables migrations for a specific collection. Uses `setAllowMigrations` admin command.',
+              example: 'sh.disableMigrations(ns)',
+            },
             getBalancerState: {
               link: 'https://mongodb.com/docs/manual/reference/method/sh.getBalancerState',
               description:

packages/shell-api/src/helpers.ts

@@ -597,6 +597,9 @@ export async function getPrintableShardStatus(
               ];
             }
 
+            collRes.allowMigrations =
+              coll.permitMigrations !== false && coll.allowMigrations !== false;
+
             const chunksRes = [];
             const chunksCollMatch = buildConfigChunksCollectionMatch(coll);
             const chunks = await (

packages/shell-api/src/shard.spec.ts

@@ -1232,6 +1232,76 @@ describe('Shard', function () {
         expect(warnSpy.calledOnce).to.equal(true);
       });
     });
+    describe('disableMigrations', function () {
+      this.beforeEach(() => {
+        serviceProvider.runCommandWithCheck.onFirstCall().resolves({
+          ok: 1,
+          msg: 'not dbgrid',
+        });
+        serviceProvider.runCommandWithCheck.onSecondCall().resolves({ ok: 1 });
+      });
+
+      it('warns if not mongos', async function () {
+        await shard.disableMigrations('ns');
+        expect(warnSpy.calledOnce).to.equal(true);
+      });
+
+      it('calls serviceProvider.runCommandWithCheck', async function () {
+        await shard.disableMigrations('ns');
+
+        expect(serviceProvider.runCommandWithCheck).to.have.been.calledWith(
+          ADMIN_DB,
+          {
+            setAllowMigrations: 'ns',
+            allowMigrations: false,
+          }
+        );
+      });
+
+      it('throws if serviceProvider.runCommandWithCheck rejects', async function () {
+        const expectedError = new Error();
+        serviceProvider.runCommandWithCheck
+          .onSecondCall()
+          .rejects(expectedError);
+        const caughtError = await shard.disableMigrations('ns').catch((e) => e);
+        expect(caughtError).to.equal(expectedError);
+      });
+    });
+    describe('enableMigrations', function () {
+      this.beforeEach(() => {
+        serviceProvider.runCommandWithCheck.onFirstCall().resolves({
+          ok: 1,
+          msg: 'not dbgrid',
+        });
+        serviceProvider.runCommandWithCheck.onSecondCall().resolves({ ok: 1 });
+      });
+
+      it('warns if not mongos', async function () {
+        await shard.enableMigrations('ns');
+        expect(warnSpy.calledOnce).to.equal(true);
+      });
+
+      it('calls serviceProvider.runCommandWithCheck', async function () {
+        await shard.enableMigrations('ns');
+
+        expect(serviceProvider.runCommandWithCheck).to.have.been.calledWith(
+          ADMIN_DB,
+          {
+            setAllowMigrations: 'ns',
+            allowMigrations: true,
+          }
+        );
+      });
+
+      it('throws if serviceProvider.runCommandWithCheck rejects', async function () {
+        const expectedError = new Error();
+        serviceProvider.runCommandWithCheck
+          .onSecondCall()
+          .rejects(expectedError);
+        const caughtError = await shard.enableMigrations('ns').catch((e) => e);
+        expect(caughtError).to.equal(expectedError);
+      });
+    });
     describe('getBalancerState', function () {
       it('returns whatever serviceProvider.find returns', async function () {
         serviceProvider.runCommandWithCheck.resolves({
@@ -2530,6 +2600,71 @@ describe('Shard', function () {
         );
       });
     });
+    describe('collection migrations', function () {
+      let db: Database;
+
+      const dbName = 'shard-status-test';
+      const ns = `${dbName}.test`;
+
+      beforeEach(async function () {
+        db = sh._database.getSiblingDB(dbName);
+        await db.getCollection('test').insertOne({ key: 1 });
+        await db.getCollection('test').createIndex({ key: 1 });
+        await sh.enableSharding(dbName);
+        await sh.shardCollection(ns, { key: 1 });
+      });
+
+      afterEach(async function () {
+        await db.dropDatabase();
+      });
+
+      const checkMigrationsEnabled = async (): Promise<boolean> => {
+        return (await sh.status()).value.databases.find(
+          (d) => d.database._id === dbName
+        )?.collections[ns].allowMigrations;
+      };
+
+      it('has migrations enabled by default', async function () {
+        expect(await checkMigrationsEnabled()).to.be.true;
+      });
+
+      it('can disable migrations', async function () {
+        expect((await sh.disableMigrations(ns)).ok).to.equal(1);
+        expect(await checkMigrationsEnabled()).to.be.false;
+      });
+
+      it('can enable migrations', async function () {
+        // Enabled by default, so disable first
+        expect((await sh.disableMigrations(ns)).ok).to.equal(1);
+        expect(await checkMigrationsEnabled()).to.be.false;
+
+        expect((await sh.enableMigrations(ns)).ok).to.equal(1);
+        expect(await checkMigrationsEnabled()).to.be.true;
+      });
+
+      it('disabling migrations is idempotent', async function () {
+        expect(await checkMigrationsEnabled()).to.be.true;
+
+        expect((await sh.disableMigrations(ns)).ok).to.equal(1);
+        expect(await checkMigrationsEnabled()).to.be.false;
+
+        // Run disable again to check idempotency
+        expect((await sh.disableMigrations(ns)).ok).to.equal(1);
+        expect(await checkMigrationsEnabled()).to.be.false;
+      });
+
+      it('enabling migrations is idempotent', async function () {
+        expect(await checkMigrationsEnabled()).to.be.true;
+
+        // Enabling when already enabled should not do anything
+        expect((await sh.enableMigrations(ns)).ok).to.equal(1);
+        expect(await checkMigrationsEnabled()).to.be.true;
+
+        // Run enable again to check idempotency
+        expect((await sh.enableMigrations(ns)).ok).to.equal(1);
+        expect(await checkMigrationsEnabled()).to.be.true;
+      });
+    });
     describe('automerge', function () {
       it('not shown if sh.status() if not explicitly enabled', async function () {
         // It might be explicitly set from 7.0

packages/shell-api/src/shard.ts

@@ -556,6 +556,39 @@ export default class Shard<
       )) as UpdateResult;
   }
 
+  private async _setAllowMigrations(
+    ns: string,
+    allowMigrations: boolean
+  ): Promise<Document> {
+    const apiCall = `${allowMigrations ? 'enable' : 'disable'}Migrations`;
+    assertArgsDefinedType([ns], ['string'], `Shard.${apiCall}`);
+    this._emitShardApiCall(apiCall, { ns });
+
+    const helloResult = await this._database._maybeCachedHello();
+    if (helloResult.msg !== 'isdbgrid') {
+      await this._database._instanceState.printWarning(
+        'MongoshWarning: [SHAPI-10003] You are not connected to a mongos. This command may not work as expected.'
+      );
+    }
+
+    return await this._database._runAdminCommand({
+      setAllowMigrations: ns,
+      allowMigrations,
+    });
+  }
+
+  @returnsPromise
+  @apiVersions([])
+  async enableMigrations(ns: string): Promise<Document> {
+    return await this._setAllowMigrations(ns, true);
+  }
+
+  @returnsPromise
+  @apiVersions([])
+  async disableMigrations(ns: string): Promise<Document> {
+    return await this._setAllowMigrations(ns, false);
+  }
+
   @returnsPromise
   @apiVersions([])
   async getBalancerState(): Promise<boolean> {