Skip to content

Commit 235c047

Browse files
authored
chore(shell-api): unified shard distribution helper MONGOSH-1694 (#2098)
* unified shard helper * allow user option to pass through * fix i8n test
1 parent 69f8e9c commit 235c047

File tree

3 files changed

+112
-0
lines changed

3 files changed

+112
-0
lines changed

packages/i18n/src/locales/en_US.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2045,6 +2045,12 @@ const translations: Catalog = {
20452045
'Returns a cursor with information about metadata inconsistencies',
20462046
example: 'sh.checkMetadataConsistency(<options>)',
20472047
},
2048+
shardAndDistributeCollection: {
2049+
description:
2050+
'Shards a collection and then immediately reshards the collection to the same shard key.',
2051+
example:
2052+
'sh.shardAndDistributeCollection(ns, key, unique?, options?)',
2053+
},
20482054
moveCollection: {
20492055
link: 'https://docs.mongodb.com/manual/reference/method/sh.moveCollection',
20502056
description:

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

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1791,6 +1791,85 @@ describe('Shard', function () {
17911791
});
17921792
});
17931793

1794+
describe('shardAndDistributeCollection', function () {
1795+
it('calls shardCollection and then reshardCollection with correct parameters', async function () {
1796+
const expectedResult = { ok: 1 };
1797+
1798+
const shardCollectionStub = sinon
1799+
.stub(shard, 'shardCollection')
1800+
.resolves(expectedResult);
1801+
const reshardCollectionStub = sinon
1802+
.stub(shard, 'reshardCollection')
1803+
.resolves(expectedResult);
1804+
1805+
await shard.shardAndDistributeCollection(
1806+
'db.coll',
1807+
{ key: 1 },
1808+
true,
1809+
{}
1810+
);
1811+
1812+
expect(shardCollectionStub.calledOnce).to.equal(true);
1813+
expect(shardCollectionStub.firstCall.args).to.deep.equal([
1814+
'db.coll',
1815+
{
1816+
key: 1,
1817+
},
1818+
true,
1819+
{},
1820+
]);
1821+
1822+
expect(reshardCollectionStub.calledOnce).to.equal(true);
1823+
expect(reshardCollectionStub.firstCall.args).to.deep.equal([
1824+
'db.coll',
1825+
{ key: 1 },
1826+
{ numInitialChunks: 1000, forceRedistribution: true },
1827+
]);
1828+
});
1829+
1830+
it('allows user to pass numInitialChunks', async function () {
1831+
const expectedResult = { ok: 1 };
1832+
1833+
const shardCollectionStub = sinon
1834+
.stub(shard, 'shardCollection')
1835+
.resolves(expectedResult);
1836+
const reshardCollectionStub = sinon
1837+
.stub(shard, 'reshardCollection')
1838+
.resolves(expectedResult);
1839+
1840+
await shard.shardAndDistributeCollection('db.coll', { key: 1 }, true, {
1841+
numInitialChunks: 1,
1842+
});
1843+
1844+
expect(shardCollectionStub.calledOnce).to.equal(true);
1845+
expect(shardCollectionStub.firstCall.args).to.deep.equal([
1846+
'db.coll',
1847+
{
1848+
key: 1,
1849+
},
1850+
true,
1851+
{
1852+
numInitialChunks: 1,
1853+
},
1854+
]);
1855+
1856+
expect(reshardCollectionStub.calledOnce).to.equal(true);
1857+
expect(reshardCollectionStub.firstCall.args).to.deep.equal([
1858+
'db.coll',
1859+
{ key: 1 },
1860+
{ numInitialChunks: 1, forceRedistribution: true },
1861+
]);
1862+
});
1863+
it('returns whatever shard.reshardCollection returns', async function () {
1864+
const expectedResult = { ok: 1 };
1865+
sinon.stub(shard, 'reshardCollection').resolves(expectedResult);
1866+
const result = await shard.shardAndDistributeCollection('db.coll', {
1867+
key: 1,
1868+
});
1869+
expect(result).to.deep.equal(expectedResult);
1870+
});
1871+
});
1872+
17941873
describe('moveCollection', function () {
17951874
it('calls serviceProvider.runCommandWithCheck', async function () {
17961875
await shard.moveCollection('db.coll', 'shard1');

packages/shell-api/src/shard.ts

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -694,6 +694,33 @@ export default class Shard extends ShellApiWithMongoClass {
694694
});
695695
}
696696

697+
@returnsPromise
698+
@serverVersions(['5.0.0', ServerVersions.latest])
699+
async shardAndDistributeCollection(
700+
ns: string,
701+
key: Document,
702+
unique?: boolean | Document,
703+
options?: Document
704+
): Promise<Document> {
705+
this._emitShardApiCall('shardAndDistributeCollection', {
706+
ns,
707+
key,
708+
unique,
709+
options,
710+
});
711+
await this.shardCollection(ns, key, unique, options);
712+
// SERVER-92762: Prevent unequal data distribution by setting
713+
// numInitialChunks to 1000.
714+
const numInitialChunks =
715+
typeof unique === 'object'
716+
? unique.numInitialChunks
717+
: options?.numInitialChunks;
718+
return await this.reshardCollection(ns, key, {
719+
numInitialChunks: numInitialChunks ?? 1000,
720+
forceRedistribution: true,
721+
});
722+
}
723+
697724
@serverVersions(['8.0.0', ServerVersions.latest])
698725
@apiVersions([])
699726
@returnsPromise

0 commit comments

Comments
 (0)