chore(shell-api): unified shard distribution helper MONGOSH-1694 (#2098)

* unified shard helper * allow user option to pass through * fix i8n test
mongodb-js · Jul 29, 2024 · 235c047 · 235c047
1 parent 69f8e9c
commit 235c047
Show file tree

Hide file tree

Showing 3 changed files with 112 additions and 0 deletions.
diff --git a/packages/i18n/src/locales/en_US.ts b/packages/i18n/src/locales/en_US.ts
@@ -2045,6 +2045,12 @@ const translations: Catalog = {
                 'Returns a cursor with information about metadata inconsistencies',
               example: 'sh.checkMetadataConsistency(<options>)',
             },
+            shardAndDistributeCollection: {
+              description:
+                'Shards a collection and then immediately reshards the collection to the same shard key.',
+              example:
+                'sh.shardAndDistributeCollection(ns, key, unique?, options?)',
+            },
             moveCollection: {
               link: 'https://docs.mongodb.com/manual/reference/method/sh.moveCollection',
               description:

diff --git a/packages/shell-api/src/shard.spec.ts b/packages/shell-api/src/shard.spec.ts
@@ -1791,6 +1791,85 @@ describe('Shard', function () {
       });
     });
 
+    describe('shardAndDistributeCollection', function () {
+      it('calls shardCollection and then reshardCollection with correct parameters', async function () {
+        const expectedResult = { ok: 1 };
+
+        const shardCollectionStub = sinon
+          .stub(shard, 'shardCollection')
+          .resolves(expectedResult);
+        const reshardCollectionStub = sinon
+          .stub(shard, 'reshardCollection')
+          .resolves(expectedResult);
+
+        await shard.shardAndDistributeCollection(
+          'db.coll',
+          { key: 1 },
+          true,
+          {}
+        );
+
+        expect(shardCollectionStub.calledOnce).to.equal(true);
+        expect(shardCollectionStub.firstCall.args).to.deep.equal([
+          'db.coll',
+          {
+            key: 1,
+          },
+          true,
+          {},
+        ]);
+
+        expect(reshardCollectionStub.calledOnce).to.equal(true);
+        expect(reshardCollectionStub.firstCall.args).to.deep.equal([
+          'db.coll',
+          { key: 1 },
+          { numInitialChunks: 1000, forceRedistribution: true },
+        ]);
+      });
+
+      it('allows user to pass numInitialChunks', async function () {
+        const expectedResult = { ok: 1 };
+
+        const shardCollectionStub = sinon
+          .stub(shard, 'shardCollection')
+          .resolves(expectedResult);
+        const reshardCollectionStub = sinon
+          .stub(shard, 'reshardCollection')
+          .resolves(expectedResult);
+
+        await shard.shardAndDistributeCollection('db.coll', { key: 1 }, true, {
+          numInitialChunks: 1,
+        });
+
+        expect(shardCollectionStub.calledOnce).to.equal(true);
+        expect(shardCollectionStub.firstCall.args).to.deep.equal([
+          'db.coll',
+          {
+            key: 1,
+          },
+          true,
+          {
+            numInitialChunks: 1,
+          },
+        ]);
+
+        expect(reshardCollectionStub.calledOnce).to.equal(true);
+        expect(reshardCollectionStub.firstCall.args).to.deep.equal([
+          'db.coll',
+          { key: 1 },
+          { numInitialChunks: 1, forceRedistribution: true },
+        ]);
+      });
+      it('returns whatever shard.reshardCollection returns', async function () {
+        const expectedResult = { ok: 1 };
+        sinon.stub(shard, 'reshardCollection').resolves(expectedResult);
+        const result = await shard.shardAndDistributeCollection('db.coll', {
+          key: 1,
+        });
+        expect(result).to.deep.equal(expectedResult);
+      });
+    });
+
     describe('moveCollection', function () {
       it('calls serviceProvider.runCommandWithCheck', async function () {
         await shard.moveCollection('db.coll', 'shard1');

diff --git a/packages/shell-api/src/shard.ts b/packages/shell-api/src/shard.ts
@@ -694,6 +694,33 @@ export default class Shard extends ShellApiWithMongoClass {
     });
   }
 
+  @returnsPromise
+  @serverVersions(['5.0.0', ServerVersions.latest])
+  async shardAndDistributeCollection(
+    ns: string,
+    key: Document,
+    unique?: boolean | Document,
+    options?: Document
+  ): Promise<Document> {
+    this._emitShardApiCall('shardAndDistributeCollection', {
+      ns,
+      key,
+      unique,
+      options,
+    });
+    await this.shardCollection(ns, key, unique, options);
+    // SERVER-92762: Prevent unequal data distribution by setting
+    // numInitialChunks to 1000.
+    const numInitialChunks =
+      typeof unique === 'object'
+        ? unique.numInitialChunks
+        : options?.numInitialChunks;
+    return await this.reshardCollection(ns, key, {
+      numInitialChunks: numInitialChunks ?? 1000,
+      forceRedistribution: true,
+    });
+  }
+
   @serverVersions(['8.0.0', ServerVersions.latest])
   @apiVersions([])
   @returnsPromise