feat(shell-api): added a new shell helper 'getShardedDataInformation' MONGOSH-1254 (#1334)

Author: PolPinol

packages/i18n/src/locales/en_US.ts

@@ -1635,6 +1635,10 @@ const translations: Catalog = {
               link: 'https://docs.mongodb.com/manual/reference/method/sh.setBalancerState',
               description: 'Calls sh.startBalancer if state is true, otherwise calls sh.stopBalancer',
               example: 'sh.setBalancerState(state)',
+            },
+            getShardedDataDistribution: {
+              description: 'Returns data-size distribution information for all existing sharded collections',
+              example: 'sh.getShardedDataDistribution()',
             }
           }
         }

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

@@ -4,7 +4,7 @@ import Shard from './shard';
 import { ADMIN_DB, ALL_PLATFORMS, ALL_SERVER_VERSIONS, ALL_TOPOLOGIES } from './enums';
 import { signatures, toShellResult } from './index';
 import Mongo from './mongo';
-import { bson, ServiceProvider, FindCursor as ServiceProviderCursor } from '@mongosh/service-provider-core';
+import { bson, ServiceProvider, FindCursor as ServiceProviderCursor, AggregationCursor as ServiceProviderAggCursor } from '@mongosh/service-provider-core';
 import { EventEmitter } from 'events';
 import ShellInstanceState from './shell-instance-state';
 import { UpdateResult } from './result';
@@ -1184,6 +1184,21 @@ describe('Shard', () => {
         expect(caughtError).to.equal(expectedError);
       });
     });
+    describe('getShardedDataDistribution', () => {
+      it('throws if aggregateDb fails', async() => {
+        serviceProvider.aggregateDb.throws(new Error('err'));
+        const error: any = await shard.getShardedDataDistribution().catch(err => err);
+        expect(error.message).to.be.equal('err');
+      });
+
+      it('throws if not mongos', async() => {
+        const serviceProviderCursor = stubInterface<ServiceProviderAggCursor>();
+        serviceProvider.aggregateDb.returns(serviceProviderCursor as any);
+        serviceProviderCursor.hasNext.throws(Object.assign(new Error(), { code: 40324 }));
+        const error: any = await shard.getShardedDataDistribution().catch(err => err);
+        expect(error.message).to.match(/sh\.getShardedDataDistribution only works on mongos/);
+      });
+    });
   });
 
   describe('integration', () => {

packages/shell-api/src/shard.ts

@@ -1,15 +1,17 @@
 import Database from './database';
 import {
   shellApiClassDefault,
-  returnsPromise, serverVersions, apiVersions, ShellApiWithMongoClass
+  returnsPromise, serverVersions, apiVersions, ShellApiWithMongoClass, returnType
 } from './decorators';
 
 import type { Document } from '@mongosh/service-provider-core';
 import { assertArgsDefinedType, getConfigDB, getPrintableShardStatus } from './helpers';
 import { ServerVersions, asPrintable } from './enums';
 import { CommandResult, UpdateResult } from './result';
 import { redactURICredentials } from '@mongosh/history';
+import { CommonErrors, MongoshRuntimeError } from '@mongosh/errors';
 import Mongo from './mongo';
+import AggregationCursor from './aggregation-cursor';
 
 @shellApiClassDefault
 export default class Shard extends ShellApiWithMongoClass {
@@ -424,4 +426,29 @@ export default class Shard extends ShellApiWithMongoClass {
     }
     return this.stopBalancer();
   }
+
+  @returnsPromise
+  @apiVersions([])
+  @serverVersions(['6.1.0', ServerVersions.latest])
+  @returnType('AggregationCursor')
+  async getShardedDataDistribution(options = {}): Promise<AggregationCursor> {
+    this._emitShardApiCall('getShardedDataDistribution', {});
+
+    const cursor = await this._database.getSiblingDB('admin').aggregate([{ $shardedDataDistribution: options }]);
+
+    try {
+      await cursor.hasNext();
+    } catch (err: any) {
+      if (err.code?.valueOf() === 40324) { // unrecognized stage error
+        throw new MongoshRuntimeError(
+          'sh.getShardedDataDistribution only works on mongos',
+          CommonErrors.CommandFailed
+        );
+      }
+
+      throw err;
+    }
+
+    return cursor;
+  }
 }