feat(shell-api): account for unsharded collections in sharding catalog COMPASS-7602 (#1808)

Author: addaleax

packages/shell-api/src/collection.ts

@@ -40,6 +40,7 @@ import {
   shouldRunAggregationImmediately,
   coerceToJSNumber,
   buildConfigChunksCollectionMatch,
+  onlyShardedCollectionsInConfigFilter,
 } from './helpers';
 import type {
   AnyBulkWriteOperation,
@@ -1762,9 +1763,7 @@ export default class Collection extends ShellApiWithMongoClass {
     try {
       result.sharded = !!(await config.getCollection('collections').findOne({
         _id: timeseriesBucketsNs ?? ns,
-        // Dropped is gone on newer server versions, so check for !== true
-        // rather than for === false (SERVER-51880 and related).
-        dropped: { $ne: true },
+        ...onlyShardedCollectionsInConfigFilter,
       }));
     } catch (e) {
       // A user might not have permissions to check the config. In which
@@ -2068,9 +2067,7 @@ export default class Collection extends ShellApiWithMongoClass {
       .getCollection('collections')
       .findOne({
         _id: ns,
-        // dropped is gone on newer server versions, so check for !== true
-        // rather than for === false (SERVER-51880 and related)
-        dropped: { $ne: true },
+        ...onlyShardedCollectionsInConfigFilter,
       });
     if (!configCollectionsInfo) {
       throw new MongoshInvalidInputError(

packages/shell-api/src/helpers.ts

@@ -27,7 +27,7 @@ import type {
   bson,
 } from '@mongosh/service-provider-core';
 import type { ClientSideFieldLevelEncryptionOptions } from './field-level-encryption';
-import type { AutoEncryptionOptions } from 'mongodb';
+import { type AutoEncryptionOptions } from 'mongodb';
 import { shellApiType } from './enums';
 import type { AbstractCursor } from './abstract-cursor';
 import type ChangeStreamCursor from './change-stream-cursor';
@@ -92,6 +92,17 @@ function getAssertCaller(caller?: string): string {
   return caller ? ` (${caller})` : '';
 }
 
+// Fields to add to a $match/filter on config.collections to only get
+// collections that are actually sharded
+export const onlyShardedCollectionsInConfigFilter = {
+  // dropped is gone on newer server versions, so check for !== true
+  // rather than for === false (SERVER-51880 and related)
+  dropped: { $ne: true },
+  // unsplittable introduced in PM-3364 to mark unsharded collections
+  // that are still being tracked in the catalog
+  unsplittable: { $ne: true },
+} as const;
+
 export function assertArgsDefinedType(
   args: any[],
   expectedTypes: Array<true | string | Array<string | undefined>>,
@@ -487,33 +498,45 @@ export async function getPrintableShardStatus(
   ]);
   result.balancer = balancerRes;
 
-  const databases = await (await configDB.getCollection('databases').find())
-    .sort({ name: 1 })
-    .toArray();
-
+  // All databases in config.databases + those implicitly referenced
+  // by a sharded collection in config.collections
+  // (could become a single pipeline using $unionWith when we drop 4.2 server support)
+  const [databases, collections] = await Promise.all([
+    (async () =>
+      await (await configDB.getCollection('databases').find())
+        .sort({ _id: 1 })
+        .toArray())(),
+    (async () =>
+      await (
+        await configDB.getCollection('collections').find({
+          ...onlyShardedCollectionsInConfigFilter,
+        })
+      )
+        .sort({ _id: 1 })
+        .toArray())(),
+  ]);
   // Special case the config db, since it doesn't have a record in config.databases.
   databases.push({ _id: 'config', primary: 'config', partitioned: true });
+
+  for (const coll of collections) {
+    if (!databases.find((db) => coll._id.startsWith(db._id + '.'))) {
+      databases.push({ _id: coll._id.split('.')[0] });
+    }
+  }
+
   databases.sort((a: any, b: any): number => {
     return a._id.localeCompare(b._id);
   });
 
-  result.databases = await Promise.all(
-    databases.map(async (db) => {
-      const escapeRegex = (string: string): string => {
-        return string.replace(/[-/\\^$*+?.()|[\]{}]/g, '\\$&');
-      };
-      const colls = await (
-        await configDB
-          .getCollection('collections')
-          .find({ _id: new RegExp('^' + escapeRegex(db._id) + '\\.') })
-      )
-        .sort({ _id: 1 })
-        .toArray();
+  result.databases = (
+    await Promise.all(
+      databases.map(async (db) => {
+        const colls = collections.filter((coll) =>
+          coll._id.startsWith(db._id + '.')
+        );
 
-      const collList: any = await Promise.all(
-        colls
-          .filter((coll) => !coll.dropped)
-          .map(async (coll) => {
+        const collList = await Promise.all(
+          colls.map(async (coll) => {
             const collRes = {} as any;
             collRes.shardKey = coll.key;
             collRes.unique = !!coll.unique;
@@ -605,12 +628,13 @@ export async function getPrintableShardStatus(
             }
             collRes.chunks = chunksRes;
             collRes.tags = tagsRes;
-            return [coll._id, collRes];
+            return [coll._id, collRes] as const;
           })
-      );
-      return { database: db, collections: Object.fromEntries(collList) };
-    })
-  );
+        );
+        return { database: db, collections: Object.fromEntries(collList) };
+      })
+    )
+  ).filter((dbEntry) => !!dbEntry);
 
   delete result.shardingVersion.currentVersion;
   return result;

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

@@ -1835,6 +1835,7 @@ describe('Shard', function () {
         .toArray();
       expect(members.length).to.equal(2);
       await sh._database.getSiblingDB(dbName).dropDatabase();
+      await sh._database.getSiblingDB(dbName).createCollection('unsharded');
     });
 
     after(function () {
@@ -1901,19 +1902,25 @@ describe('Shard', function () {
     });
     describe('turn on sharding', function () {
       it('enableSharding for a db', async function () {
-        expect((await sh.status()).value.databases.length).to.equal(1);
-        expect((await sh.enableSharding(dbName)).ok).to.equal(1);
-        expect((await sh.status()).value.databases.length).to.equal(2);
+        expect((await sh.status()).value.databases.length).to.oneOf([1, 2]);
+        expect((await sh.enableSharding(dbName)).ok).to.equal(1); // This may not have any effect on newer server versions
+        expect((await sh.status()).value.databases.length).to.be.oneOf([1, 2]);
       });
       it('enableSharding for a collection and modify documents in it', async function () {
         expect(
-          Object.keys((await sh.status()).value.databases[1].collections).length
-        ).to.equal(0);
+          Object.keys(
+            (await sh.status()).value.databases.find(
+              (d) => d.database._id === 'test'
+            )?.collections ?? []
+          )
+        ).to.deep.equal([]);
         expect(
           (await sh.shardCollection(ns, { key: 1 })).collectionsharded
         ).to.equal(ns);
         expect(
-          (await sh.status()).value.databases[1].collections[ns].shardKey
+          (await sh.status()).value.databases.find(
+            (d) => d.database._id === 'test'
+          ).collections[ns].shardKey
         ).to.deep.equal({ key: 1 });
 
         const db = instanceState.currentDb.getSiblingDB(dbName);
@@ -2122,6 +2129,7 @@ describe('Shard', function () {
       });
       it('fails when running against an unsharded collection', async function () {
         try {
+          await db.createCollection('test');
           await db.getCollection('test').getShardDistribution();
         } catch (err: any) {
           expect(err.name).to.equal('MongoshInvalidInputError');