update collection.spec.ts, collection.ts and result.ts

mongodb-js · Oct 8, 2024 · 59e6f25 · 59e6f25
1 parent 37b7da1
commit 59e6f25
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 4 deletions.
diff --git 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<ServiceProviderCursor>();
+
+          // 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
@@ -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<CommandResult> {
+  async getShardDistribution(): Promise<
+    CommandResult<GetShardDistributionResult>
+  > {
     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<GetShardDistributionResult>('StatsResult', result);
   }
 
   @serverVersions(['3.1.0', ServerVersions.latest])

diff --git 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;
+  };
+};