diff --git a/packages/shell-api/src/collection.spec.ts b/packages/shell-api/src/collection.spec.ts index 205da41f3..f29729eec 100644 --- a/packages/shell-api/src/collection.spec.ts +++ b/packages/shell-api/src/collection.spec.ts @@ -2276,6 +2276,70 @@ describe('Collection', function () { ShellApiErrors.NotConnectedToShardedCluster ); }); + + describe('with orphan documents', function () { + const mockedNumChunks = 2; + const mockedCollectionConfigInfo = {}; + const mockedShardStats = { + shard: 'test-shard', + storageStats: { + size: 1000, + numOrphanDocs: 10, + avgObjSize: 7, + count: 15, + }, + }; + const mockedShardInfo = { + host: 'dummy-host', + }; + + beforeEach(function () { + const serviceProviderCursor = stubInterface(); + + // Make find and limit have no effect so the value of findOne is determined by tryNext. + serviceProviderCursor.limit.returns(serviceProviderCursor); + serviceProvider.find.returns(serviceProviderCursor); + + // Mock according to the order of findOne calls getShardDistribution uses. + serviceProviderCursor.tryNext + .onCall(0) + .resolves(mockedCollectionConfigInfo); + serviceProviderCursor.tryNext.onCall(1).resolves(mockedShardInfo); + serviceProvider.countDocuments.returns( + Promise.resolve(mockedNumChunks) + ); + + const aggregateTryNext = sinon.stub(); + aggregateTryNext.onCall(0).resolves(mockedShardStats); + aggregateTryNext.onCall(1).resolves(null); + + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + serviceProvider.aggregate.returns({ + tryNext: aggregateTryNext, + } as any); + }); + + it('should account for numOrphanDocs when calculating size', async function () { + const shardDistribution = await collection.getShardDistribution(); + + const { storageStats } = mockedShardStats; + expect(shardDistribution.type).equals('StatsResult'); + const adjustedSize = + storageStats.size - + storageStats.numOrphanDocs * storageStats.avgObjSize; + expect(shardDistribution.value.Totals.data).equals( + `${adjustedSize}B` + ); + const shardField = Object.keys(shardDistribution.value).find( + (field) => field !== 'Totals' + ) as string; + + expect(shardField).not.undefined; + expect( + shardDistribution.value[shardField]['estimated data per chunk'] + ).equals(`${adjustedSize / mockedNumChunks}B`); + }); + }); }); describe('analyzeShardKey', function () { diff --git a/packages/shell-api/src/collection.ts b/packages/shell-api/src/collection.ts index 51f85a846..ee16a22bc 100644 --- a/packages/shell-api/src/collection.ts +++ b/packages/shell-api/src/collection.ts @@ -90,6 +90,7 @@ import { HIDDEN_COMMANDS } from '@mongosh/history'; import PlanCache from './plan-cache'; import ChangeStreamCursor from './change-stream-cursor'; import { ShellApiErrors } from './error-codes'; +import type { GetShardDistributionResult } from './result'; @shellApiClassDefault @addSourceToResults @@ -2135,12 +2136,14 @@ export default class Collection extends ShellApiWithMongoClass { @returnsPromise @topologies([Topologies.Sharded]) @apiVersions([]) - async getShardDistribution(): Promise { + async getShardDistribution(): Promise< + CommandResult + > { this._emitCollectionApiCall('getShardDistribution', {}); await getConfigDB(this._database); // Warns if not connected to mongos - const result = {} as Document; + const result = {} as GetShardDistributionResult; const config = this._mongo.getDB('config'); const collStats = await ( @@ -2235,7 +2238,7 @@ export default class Collection extends ShellApiWithMongoClass { data: dataFormat(totals.size), docs: totals.count, chunks: totals.numChunks, - } as Document; + } as GetShardDistributionResult['Totals']; for (const shardStats of conciseShardsStats) { const estDataPercent = @@ -2254,7 +2257,7 @@ export default class Collection extends ShellApiWithMongoClass { ]; } result.Totals = totalValue; - return new CommandResult('StatsResult', result); + return new CommandResult('StatsResult', result); } @serverVersions(['3.1.0', ServerVersions.latest]) diff --git a/packages/shell-api/src/result.ts b/packages/shell-api/src/result.ts index f4d2b3162..9500a6d32 100644 --- a/packages/shell-api/src/result.ts +++ b/packages/shell-api/src/result.ts @@ -127,3 +127,21 @@ export class CursorIterationResult extends ShellApiValueClass { this.documents = []; } } + +export type GetShardDistributionResult = { + Totals: { + data: string; + docs: number; + chunks: number; + } & { + [individualShardDistribution: string]: string[]; + }; +} & { + [individualShardResult: string]: { + data: string; + docs: number; + chunks: number; + 'estimated data per chunk': string; + 'estimated docs per chunk': number; + }; +};