Merge branch 'gagik/add-sharded-time-series-getshard' of github.com:m…

…ongodb-js/mongosh into gagik/get-shard-orphan-documents
mongodb-js · Oct 8, 2024 · 37b7da1 · 37b7da1
2 parents 20a81bc + bf1f643
commit 37b7da1
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 47 deletions.
diff --git a/packages/shell-api/src/collection.spec.ts b/packages/shell-api/src/collection.spec.ts
@@ -2258,8 +2258,16 @@ describe('Collection', function () {
       it('throws when collection is not sharded', async function () {
         const serviceProviderCursor = stubInterface<ServiceProviderCursor>();
         serviceProviderCursor.limit.returns(serviceProviderCursor);
-        serviceProviderCursor.tryNext.resolves(null);
-        serviceProvider.find.returns(serviceProviderCursor as any);
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+        serviceProviderCursor.tryNext.returns(null as any);
+        serviceProvider.find.returns(serviceProviderCursor);
+
+        const tryNext = sinon.stub();
+        tryNext.onCall(0).resolves({ storageStats: {} });
+        tryNext.onCall(1).resolves(null);
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+        serviceProvider.aggregate.returns({ tryNext } as any);
+
         const error = await collection.getShardDistribution().catch((e) => e);
 
         expect(error).to.be.instanceOf(MongoshInvalidInputError);

diff --git a/packages/shell-api/src/collection.ts b/packages/shell-api/src/collection.ts
@@ -2077,32 +2077,43 @@ export default class Collection extends ShellApiWithMongoClass {
   }
 
   /**
-   * Helper for getting collection info of sharded timeseries collections
-   * @returns collection info based on given collStats, null if the information is not found or
-   * if the collection is not timeseries.
+   * Helper for getting collection info for sharded collections.
+   * @throws If the collection is not sharded.
+   * @returns collection info based on given collStats.
    */
-  async _getShardedTimeseriesCollectionInfo(
+  async _getShardedCollectionInfo(
     config: Database,
     collStats: Document[]
-  ): Promise<Document | null> {
+  ): Promise<Document> {
+    const ns = `${this._database._name}.${this._name}`;
+    const existingConfigCollectionsInfo = await config
+      .getCollection('collections')
+      .findOne({
+        _id: ns,
+        ...onlyShardedCollectionsInConfigFilter,
+      });
+
+    if (existingConfigCollectionsInfo !== null) {
+      return existingConfigCollectionsInfo;
+    }
+
+    // If the collection info is not found, check if it is timeseries and use the bucket
     const timeseriesShardStats = collStats.find(
       (extractedShardStats) =>
         typeof extractedShardStats.storageStats.timeseries !== 'undefined'
     );
 
     if (!timeseriesShardStats) {
-      return null;
+      throw new MongoshInvalidInputError(
+        `Collection ${this._name} is not sharded`,
+        ShellApiErrors.NotConnectedToShardedCluster
+      );
     }
 
     const { storageStats } = timeseriesShardStats;
 
-    const timeseries: Document = storageStats['timeseries'];
-
-    const timeseriesBucketNs: string | undefined = timeseries['bucketsNs'];
-
-    if (!timeseriesBucketNs) {
-      return null;
-    }
+    const timeseries: Document = storageStats.timeseries;
+    const timeseriesBucketNs: string = timeseries.bucketsNs;
 
     const timeseriesCollectionInfo = await config
       .getCollection('collections')
@@ -2111,6 +2122,13 @@ export default class Collection extends ShellApiWithMongoClass {
         ...onlyShardedCollectionsInConfigFilter,
       });
 
+    if (!timeseriesCollectionInfo) {
+      throw new MongoshRuntimeError(
+        `Error finding collection information for ${timeseriesBucketNs}`,
+        CommonErrors.CommandFailed
+      );
+    }
+
     return timeseriesCollectionInfo;
   }
 
@@ -2139,30 +2157,10 @@ export default class Collection extends ShellApiWithMongoClass {
       avgObjSize: number;
     }[] = [];
 
-    const ns = `${this._database._name}.${this._name}`;
-    let configCollectionsInfo: Document;
-    const existingConfigCollectionsInfo = await config
-      .getCollection('collections')
-      .findOne({
-        _id: ns,
-        ...onlyShardedCollectionsInConfigFilter,
-      });
-
-    if (!existingConfigCollectionsInfo) {
-      const timeseriesCollectionInfo =
-        await this._getShardedTimeseriesCollectionInfo(config, collStats);
-
-      if (!timeseriesCollectionInfo) {
-        throw new MongoshInvalidInputError(
-          `Collection ${this._name} is not sharded`,
-          ShellApiErrors.NotConnectedToShardedCluster
-        );
-      }
-
-      configCollectionsInfo = timeseriesCollectionInfo;
-    } else {
-      configCollectionsInfo = existingConfigCollectionsInfo;
-    }
+    const configCollectionsInfo = await this._getShardedCollectionInfo(
+      config,
+      collStats
+    );
 
     await Promise.all(
       collStats.map((extractedShardStats) =>
@@ -2203,8 +2201,9 @@ export default class Collection extends ShellApiWithMongoClass {
 
           const key = `Shard ${shardStats.shardId} at ${shardStats.host}`;
 
-          // In sharded timeseries collections, count is 0.
-          const shardStatsCount = shardStats.count ?? 0;
+          // In sharded timeseries collections we do not have a count
+          // so we intentionally pass NaN as a result to the client.
+          const shardStatsCount: number = shardStats.count ?? NaN;
 
           const estimatedChunkDataPerChunk =
             shardStats.numChunks === 0

diff --git a/packages/shell-api/src/shard.spec.ts b/packages/shell-api/src/shard.spec.ts
@@ -2557,7 +2557,7 @@ describe('Shard', function () {
 
     describe('collection.getShardDistribution()', function () {
       let db: Database;
-      const dbName = 'shard-stats-test';
+      const dbName = 'get-shard-distribution-test';
       const ns = `${dbName}.test`;
 
       beforeEach(async function () {
@@ -2614,7 +2614,7 @@ describe('Shard', function () {
       context('sharded timeseries collections', function () {
         skipIfServerVersion(mongos, '< 5.1');
 
-        const timeseriesCollectionName = 'testTS';
+        const timeseriesCollectionName = 'getShardDistributionTS';
         const timeseriesNS = `${dbName}.${timeseriesCollectionName}`;
 
         beforeEach(async function () {
@@ -2647,7 +2647,7 @@ describe('Shard', function () {
 
         it('returns the correct StatsResult', async function () {
           const result = await db
-            .getCollection('testTS')
+            .getCollection(timeseriesCollectionName)
             .getShardDistribution();
           const shardDistributionValue = result.value as Document;
 
@@ -2659,12 +2659,12 @@ describe('Shard', function () {
           expect(shardFields.length).to.equal(1);
           const shardField = shardFields[0];
 
-          // Timeseries will have count 0
+          // Timeseries will have count NaN
           expect(
             shardDistributionValue[shardField]['estimated docs per chunk']
-          ).to.equal(0);
+          ).to.be.NaN;
 
-          expect(shardDistributionValue.Totals.docs).to.equal(0);
+          expect(shardDistributionValue.Totals.docs).to.be.NaN;
           expect(shardDistributionValue.Totals.chunks).to.equal(1);
         });
       });