From 93b4f894ab1991a7cee1c6d35807b151ab994689 Mon Sep 17 00:00:00 2001 From: Nikola Irinchev Date: Mon, 3 Mar 2025 22:30:02 +0100 Subject: [PATCH] feat: update sh.enable/disableBalancing to use configureCollectionBalancing --- packages/shell-api/src/shard.spec.ts | 387 ++++++++++++++++++++------- packages/shell-api/src/shard.ts | 56 ++-- 2 files changed, 328 insertions(+), 115 deletions(-) diff --git a/packages/shell-api/src/shard.spec.ts b/packages/shell-api/src/shard.spec.ts index bb857cf9a..892d2a50a 100644 --- a/packages/shell-api/src/shard.spec.ts +++ b/packages/shell-api/src/shard.spec.ts @@ -1090,61 +1090,160 @@ describe('Shard', function () { expect(caughtError).to.equal(expectedError); }); }); + describe('disableBalancing', function () { - it('calls serviceProvider.updateOne', async function () { - serviceProvider.runCommandWithCheck.resolves({ - ok: 1, - msg: 'isdbgrid', + let connectionInfoStub: sinon.SinonStub; + + this.afterEach(() => { + connectionInfoStub?.restore(); + }); + + for (const serverVersion of [ + '4.2.0', + '5.0.0', + '6.0.0', + '7.0.0', + '8.0.0', + ]) { + describe(`on server ${serverVersion}`, function () { + this.beforeEach(() => { + connectionInfoStub = sinon + .stub(instanceState, 'cachedConnectionInfo') + .returns({ + extraInfo: { + server_version: serverVersion, + uri: 'mongodb://localhost:27017', + }, + buildInfo: null, + }); + + serviceProvider.runCommandWithCheck.resolves({ + ok: 1, + msg: 'isdbgrid', + }); + }); + + it('calls serviceProvider.updateOne', async function () { + serviceProvider.runCommandWithCheck.resolves({ + ok: 1, + msg: 'isdbgrid', + }); + const expectedResult = { + matchedCount: 1, + modifiedCount: 1, + upsertedCount: 1, + upsertedId: { _id: 0 }, + result: { ok: 1, n: 1, nModified: 1 }, + connection: null, + } as any; + serviceProvider.updateOne.resolves(expectedResult); + await shard.disableBalancing('ns'); + + expect(serviceProvider.updateOne).to.have.been.calledWith( + 'config', + 'collections', + { _id: 'ns' }, + { $set: { noBalance: true } }, + { writeConcern: { w: 'majority', wtimeout: 60000 } } + ); + }); + + it('returns whatever serviceProvider.updateOne returns', async function () { + serviceProvider.runCommandWithCheck.resolves({ + ok: 1, + msg: 'isdbgrid', + }); + const oid = new bson.ObjectId(); + const expectedResult = { + matchedCount: 1, + modifiedCount: 1, + upsertedCount: 1, + upsertedId: oid, + result: { ok: 1, n: 1, nModified: 1 }, + connection: null, + acknowledged: true, + } as any; + serviceProvider.updateOne.resolves(expectedResult); + const result = await shard.disableBalancing('ns'); + expect(result).to.deep.equal(new UpdateResult(true, 1, 1, 1, oid)); + }); + + it('throws if serviceProvider.updateOne rejects', async function () { + serviceProvider.runCommandWithCheck.resolves({ + ok: 1, + msg: 'isdbgrid', + }); + const expectedError = new Error(); + serviceProvider.updateOne.rejects(expectedError); + const caughtError = await shard + .disableBalancing('ns') + .catch((e) => e); + expect(caughtError).to.equal(expectedError); + }); }); - const expectedResult = { - matchedCount: 1, - modifiedCount: 1, - upsertedCount: 1, - upsertedId: { _id: 0 }, - result: { ok: 1, n: 1, nModified: 1 }, - connection: null, - } as any; - serviceProvider.updateOne.resolves(expectedResult); - await shard.disableBalancing('ns'); + } - expect(serviceProvider.updateOne).to.have.been.calledWith( - 'config', - 'collections', - { _id: 'ns' }, - { $set: { noBalance: true } }, - { writeConcern: { w: 'majority', wtimeout: 60000 } } - ); - }); + describe('on server >= 8.1', function () { + this.beforeEach(() => { + connectionInfoStub = sinon + .stub(instanceState, 'cachedConnectionInfo') + .returns({ + extraInfo: { + server_version: '8.1.0', + uri: 'mongodb://localhost:27017', + }, + buildInfo: null, + }); - it('returns whatever serviceProvider.updateOne returns', async function () { - serviceProvider.runCommandWithCheck.resolves({ - ok: 1, - msg: 'isdbgrid', + serviceProvider.runCommandWithCheck.resolves({ + ok: 1, + msg: 'isdbgrid', + }); }); - const oid = new bson.ObjectId(); - const expectedResult = { - matchedCount: 1, - modifiedCount: 1, - upsertedCount: 1, - upsertedId: oid, - result: { ok: 1, n: 1, nModified: 1 }, - connection: null, - acknowledged: true, - } as any; - serviceProvider.updateOne.resolves(expectedResult); - const result = await shard.disableBalancing('ns'); - expect(result).to.deep.equal(new UpdateResult(true, 1, 1, 1, oid)); - }); - it('throws if serviceProvider.updateOne rejects', async function () { - serviceProvider.runCommandWithCheck.resolves({ - ok: 1, - msg: 'isdbgrid', + it('runs configureCollectionBalancing admin command', async function () { + await shard.disableBalancing('ns'); + + expect(serviceProvider.runCommandWithCheck).to.have.been.calledWith( + 'admin', + { + configureCollectionBalancing: 'ns', + enableBalancing: false, + noBalance: true, + } + ); + }); + + it('returns whatever serviceProvider.runCommandWithCheck returns', async function () { + const expectedResult = { + ok: 1, + $clusterTime: { + clusterTime: { $timestamp: { t: 1741010129, i: 1 } }, + signature: { + hash: { + $binary: { + base64: 'AAAAAAAAAAAAAAAAAAAAAAAAAAA=', + subType: 0, + }, + }, + keyId: { $numberLong: '0' }, + }, + }, + operationTime: { t: 1741010129, i: 1 }, + }; + serviceProvider.runCommandWithCheck.resolves(expectedResult); + const result = await shard.disableBalancing('ns'); + expect(result).to.deep.equal(expectedResult); + }); + + it('throws if serviceProvider.runCommandWithCheck rejects', async function () { + const expectedError = new Error(); + serviceProvider.runCommandWithCheck.rejects(expectedError); + const caughtError = await shard + .disableBalancing('ns') + .catch((e) => e); + expect(caughtError).to.equal(expectedError); }); - const expectedError = new Error(); - serviceProvider.updateOne.rejects(expectedError); - const caughtError = await shard.disableBalancing('ns').catch((e) => e); - expect(caughtError).to.equal(expectedError); }); it('throws if not mongos', async function () { @@ -1158,61 +1257,146 @@ describe('Shard', function () { expect(warnSpy.calledOnce).to.equal(true); }); }); + describe('enableBalancing', function () { - it('calls serviceProvider.updateOne', async function () { - serviceProvider.runCommandWithCheck.resolves({ - ok: 1, - msg: 'isdbgrid', + let connectionInfoStub: sinon.SinonStub; + + this.afterEach(() => { + connectionInfoStub?.restore(); + }); + + for (const serverVersion of [ + '4.2.0', + '5.0.0', + '6.0.0', + '7.0.0', + '8.0.0', + ]) { + describe(`on server ${serverVersion}`, function () { + this.beforeEach(() => { + connectionInfoStub = sinon + .stub(instanceState, 'cachedConnectionInfo') + .returns({ + extraInfo: { + server_version: serverVersion, + uri: 'mongodb://localhost:27017', + }, + buildInfo: null, + }); + + serviceProvider.runCommandWithCheck.resolves({ + ok: 1, + msg: 'isdbgrid', + }); + }); + + it('calls serviceProvider.updateOne', async function () { + const expectedResult = { + matchedCount: 1, + modifiedCount: 1, + upsertedCount: 1, + upsertedId: { _id: 0 }, + result: { ok: 1, n: 1, nModified: 1 }, + connection: null, + } as any; + serviceProvider.updateOne.resolves(expectedResult); + await shard.enableBalancing('ns'); + + expect(serviceProvider.updateOne).to.have.been.calledWith( + 'config', + 'collections', + { _id: 'ns' }, + { $set: { noBalance: false } }, + { writeConcern: { w: 'majority', wtimeout: 60000 } } + ); + }); + + it('returns whatever serviceProvider.updateOne returns', async function () { + const oid = new bson.ObjectId(); + const expectedResult = { + matchedCount: 1, + modifiedCount: 1, + upsertedCount: 1, + upsertedId: oid, + result: { ok: 1, n: 1, nModified: 1 }, + connection: null, + acknowledged: true, + } as any; + serviceProvider.updateOne.resolves(expectedResult); + const result = await shard.enableBalancing('ns'); + expect(result).to.deep.equal(new UpdateResult(true, 1, 1, 1, oid)); + }); + + it('throws if serviceProvider.updateOne rejects', async function () { + const expectedError = new Error(); + serviceProvider.updateOne.rejects(expectedError); + const caughtError = await shard + .enableBalancing('ns') + .catch((e) => e); + expect(caughtError).to.equal(expectedError); + }); }); - const expectedResult = { - matchedCount: 1, - modifiedCount: 1, - upsertedCount: 1, - upsertedId: { _id: 0 }, - result: { ok: 1, n: 1, nModified: 1 }, - connection: null, - } as any; - serviceProvider.updateOne.resolves(expectedResult); - await shard.enableBalancing('ns'); + } - expect(serviceProvider.updateOne).to.have.been.calledWith( - 'config', - 'collections', - { _id: 'ns' }, - { $set: { noBalance: false } }, - { writeConcern: { w: 'majority', wtimeout: 60000 } } - ); - }); + describe('on server >= 8.1', function () { + this.beforeEach(() => { + connectionInfoStub = sinon + .stub(instanceState, 'cachedConnectionInfo') + .returns({ + extraInfo: { + server_version: '8.1.0', + uri: 'mongodb://localhost:27017', + }, + buildInfo: null, + }); - it('returns whatever serviceProvider.updateOne returns', async function () { - serviceProvider.runCommandWithCheck.resolves({ - ok: 1, - msg: 'isdbgrid', + serviceProvider.runCommandWithCheck.resolves({ + ok: 1, + msg: 'isdbgrid', + }); }); - const oid = new bson.ObjectId(); - const expectedResult = { - matchedCount: 1, - modifiedCount: 1, - upsertedCount: 1, - upsertedId: oid, - result: { ok: 1, n: 1, nModified: 1 }, - connection: null, - acknowledged: true, - } as any; - serviceProvider.updateOne.resolves(expectedResult); - const result = await shard.enableBalancing('ns'); - expect(result).to.deep.equal(new UpdateResult(true, 1, 1, 1, oid)); - }); - it('throws if serviceProvider.updateOne rejects', async function () { - serviceProvider.runCommandWithCheck.resolves({ - ok: 1, - msg: 'isdbgrid', + it('runs configureCollectionBalancing admin command', async function () { + await shard.enableBalancing('ns'); + + expect(serviceProvider.runCommandWithCheck).to.have.been.calledWith( + 'admin', + { + configureCollectionBalancing: 'ns', + enableBalancing: true, + noBalance: false, + } + ); + }); + + it('returns whatever serviceProvider.runCommandWithCheck returns', async function () { + const expectedResult = { + ok: 1, + $clusterTime: { + clusterTime: { $timestamp: { t: 1741010129, i: 1 } }, + signature: { + hash: { + $binary: { + base64: 'AAAAAAAAAAAAAAAAAAAAAAAAAAA=', + subType: 0, + }, + }, + keyId: { $numberLong: '0' }, + }, + }, + operationTime: { t: 1741010129, i: 1 }, + }; + serviceProvider.runCommandWithCheck.resolves(expectedResult); + const result = await shard.enableBalancing('ns'); + expect(result).to.deep.equal(expectedResult); + }); + + it('throws if serviceProvider.runCommandWithCheck rejects', async function () { + const expectedError = new Error(); + serviceProvider.runCommandWithCheck.rejects(expectedError); + const caughtError = await shard.enableBalancing('ns').catch((e) => e); + expect(caughtError).to.equal(expectedError); }); - const expectedError = new Error(); - serviceProvider.updateOne.rejects(expectedError); - const caughtError = await shard.enableBalancing('ns').catch((e) => e); - expect(caughtError).to.equal(expectedError); }); it('throws if not mongos', async function () { @@ -1226,6 +1410,7 @@ describe('Shard', function () { expect(warnSpy.calledOnce).to.equal(true); }); }); + describe('getBalancerState', function () { it('returns whatever serviceProvider.find returns', async function () { serviceProvider.runCommandWithCheck.resolves({ @@ -2240,6 +2425,7 @@ describe('Shard', function () { }); }); }); + describe('turn on sharding', function () { it('enableSharding for a db', async function () { expect((await sh.status()).value.databases.length).to.oneOf([1, 2]); @@ -2287,6 +2473,7 @@ describe('Shard', function () { ); }); }); + describe('autosplit', function () { skipIfServerVersion(mongos, '> 6.x'); // Auto-splitter is removed in 7.0 it('disables correctly', async function () { @@ -2302,6 +2489,7 @@ describe('Shard', function () { ).to.equal('yes'); }); }); + describe('tags', function () { it('creates a zone', async function () { expect((await sh.addShardTag(`${shardId}-1`, 'zone1')).ok).to.equal(1); @@ -2412,6 +2600,7 @@ describe('Shard', function () { } }); }); + describe('balancer', function () { it('reports balancer state', async function () { expect(Object.keys(await sh.isBalancerRunning())).to.include.members([ @@ -2459,6 +2648,7 @@ describe('Shard', function () { ).to.equal(false); }); }); + describe('autoMerger', function () { it('reports autoMerger state', async function () { expect(await sh.isAutoMergerEnabled()).to.equal(true); @@ -2503,6 +2693,7 @@ describe('Shard', function () { ).to.not.exist; }); }); + describe('getShardDistribution', function () { let db: Database; const dbName = 'shard-distrib-test'; @@ -2595,6 +2786,7 @@ describe('Shard', function () { } }); }); + describe('analyzeShardKey()', function () { skipIfServerVersion(mongos, '< 7.0'); // analyzeShardKey will only be added in 7.0 which is not included in stable yet @@ -2630,6 +2822,7 @@ describe('Shard', function () { ).to.deep.include({ ok: 1 }); }); }); + describe('configureQueryAnalyzer()', function () { skipIfServerVersion(mongos, '< 7.0'); // analyzeShardKey will only be added in 7.0 which is not included in stable yet @@ -3034,6 +3227,7 @@ describe('Shard', function () { expect(ret).to.equal(true); }); }); + describe('databases', function () { let dbRegular: Database; let dbSh: Database; @@ -3084,6 +3278,7 @@ describe('Shard', function () { ]); }); }); + describe('checkMetadataConsistency', function () { skipIfServerVersion(mongos, '< 7.0'); let db: Database; diff --git a/packages/shell-api/src/shard.ts b/packages/shell-api/src/shard.ts index bf9dda12e..772e124bc 100644 --- a/packages/shell-api/src/shard.ts +++ b/packages/shell-api/src/shard.ts @@ -492,34 +492,52 @@ export default class Shard extends ShellApiWithMongoClass { }); } - @returnsPromise - @apiVersions([]) - async enableBalancing(ns: string): Promise { - assertArgsDefinedType([ns], ['string'], 'Shard.enableBalancing'); - this._emitShardApiCall('enableBalancing', { ns }); + private async updateBalancing( + ns: string, + enabled: boolean + ): Promise { + const apiCall = enabled ? 'enableBalancing' : 'disableBalancing'; + assertArgsDefinedType([ns], ['string'], `Shard.${apiCall}`); + this._emitShardApiCall(apiCall, { ns }); + + const serverVersion = ( + this._instanceState.cachedConnectionInfo() ?? + (await this._instanceState.fetchConnectionInfo()) + )?.extraInfo?.server_version; + + if (serverVersion && semver.gte(serverVersion, '8.1.0-alpha3')) { + // 8.1 and later support enabling/disabling balancing using the configureCollectionBalancing command. + // This is preferred over manually updating the config database, because it has additional validations + // and works for Atlas Clusters, where updating the config namespaces is not allowed. See + // https://jira.mongodb.org/browse/MONGOSH-1932 and https://jira.mongodb.org/browse/SERVER-78849 for + // more details. + return await this._database._runAdminCommand({ + configureCollectionBalancing: ns, + noBalance: !enabled, + enableBalancing: enabled, // This will be the new name after SERVER-100540 gets implemented + }); + } + const config = await getConfigDB(this._database); - return (await config + return await config .getCollection('collections') .updateOne( { _id: ns }, - { $set: { noBalance: false } }, + { $set: { noBalance: !enabled } }, { writeConcern: { w: 'majority', wtimeout: 60000 } } - )) as UpdateResult; + ); } @returnsPromise @apiVersions([]) - async disableBalancing(ns: string): Promise { - assertArgsDefinedType([ns], ['string'], 'Shard.disableBalancing'); - this._emitShardApiCall('disableBalancing', { ns }); - const config = await getConfigDB(this._database); - return (await config - .getCollection('collections') - .updateOne( - { _id: ns }, - { $set: { noBalance: true } }, - { writeConcern: { w: 'majority', wtimeout: 60000 } } - )) as UpdateResult; + enableBalancing(ns: string): Promise { + return this.updateBalancing(ns, true); + } + + @returnsPromise + @apiVersions([]) + disableBalancing(ns: string): Promise { + return this.updateBalancing(ns, false); } @returnsPromise