feat(autocomplete): autocomplete collection names MONGOSH-335

Also implicitly address MONGOSH-291 by adding a 200ms timeout after
which cached results are used (which may seem either high or low,
so I’m happy to adjust it – see also the comment about that value).

Author: addaleax

packages/autocomplete/index.spec.ts

@@ -3,21 +3,24 @@ import { signatures as shellSignatures, Topologies } from '@mongosh/shell-api';
 
 import { expect } from 'chai';
 
+let collections: string[];
 const standalone440 = {
   topology: () => Topologies.Standalone,
   connectionInfo: () => ({
     is_atlas: false,
     is_data_lake: false,
     server_version: '4.4.0'
-  })
+  }),
+  getCollectionCompletionsForCurrentDb: () => collections
 };
 const sharded440 = {
   topology: () => Topologies.Sharded,
   connectionInfo: () => ({
     is_atlas: false,
     is_data_lake: false,
     server_version: '4.4.0'
-  })
+  }),
+  getCollectionCompletionsForCurrentDb: () => collections
 };
 
 const standalone300 = {
@@ -26,23 +29,30 @@ const standalone300 = {
     is_atlas: false,
     is_data_lake: false,
     server_version: '3.0.0'
-  })
+  }),
+  getCollectionCompletionsForCurrentDb: () => collections
 };
 const datalake440 = {
   topology: () => Topologies.Sharded,
   connectionInfo: () => ({
     is_atlas: true,
     is_data_lake: true,
     server_version: '4.4.0'
-  })
+  }),
+  getCollectionCompletionsForCurrentDb: () => collections
 };
 
 const noParams = {
   topology: () => Topologies.Standalone,
-  connectionInfo: () => undefined
+  connectionInfo: () => undefined,
+  getCollectionCompletionsForCurrentDb: () => collections
 };
 
 describe('completer.completer', () => {
+  beforeEach(() => {
+    collections = [];
+  });
+
   context('when context is top level shell api', () => {
     it('matches shell completions', async() => {
       const i = 'u';
@@ -189,6 +199,12 @@ describe('completer.completer', () => {
       const i = 'db.shipw';
       expect(await completer(standalone440, i)).to.deep.equal([[], i]);
     });
+
+    it('includes collection names', async() => {
+      collections = ['shipwrecks'];
+      const i = 'db.shipw';
+      expect(await completer(standalone440, i)).to.deep.equal([['db.shipwrecks'], i]);
+    });
   });
 
   context('when context is collections', () => {

packages/autocomplete/index.ts

@@ -20,7 +20,8 @@ export interface AutocompleteParameters {
     is_atlas: boolean;
     is_data_lake: boolean;
     server_version: string;
-  };
+  },
+  getCollectionCompletionsForCurrentDb: (collName: string) => string[] | Promise<string[]>;
 }
 
 export const BASE_COMPLETIONS = EXPRESSION_OPERATORS.concat(
@@ -67,6 +68,8 @@ async function completer(params: AutocompleteParameters, line: string): Promise<
     return [hits.length ? hits : [], line];
   } else if (firstLineEl.includes('db') && splitLine.length === 2) {
     const hits = filterShellAPI(params, DB_COMPLETIONS, elToComplete, splitLine);
+    const colls = await params.getCollectionCompletionsForCurrentDb(elToComplete.trim());
+    hits.push(...colls.map(coll => `${splitLine[0]}.${coll}`));
     return [hits.length ? hits : [], line];
   } else if (firstLineEl.includes('db') && splitLine.length > 2) {
     if (splitLine.length > 3) {

packages/browser-runtime-core/src/autocompleter/shell-api-autocompleter.spec.ts

@@ -8,7 +8,8 @@ const standalone440 = {
     is_atlas: false,
     is_data_lake: false,
     server_version: '4.4.0'
-  })
+  }),
+  getCollectionCompletionsForCurrentDb: () => ['bananas']
 };
 
 describe('Autocompleter', () => {
@@ -34,6 +35,14 @@ describe('Autocompleter', () => {
         completion: 'db.coll1.find'
       });
     });
+
+    it('returns collection names value with text after dot', async() => {
+      const completions = await autocompleter.getCompletions('db.b');
+
+      expect(completions).to.deep.contain({
+        completion: 'db.bananas'
+      });
+    });
   });
 });
 

packages/cli-repl/src/cli-repl.spec.ts

@@ -5,7 +5,7 @@ import { promisify } from 'util';
 import { once } from 'events';
 import CliRepl, { CliReplOptions } from './cli-repl';
 import { startTestServer, MongodSetup } from '../../../testing/integration-testing-hooks';
-import { expect, useTmpdir, waitEval, waitBus, fakeTTYProps, readReplLogfile, tick } from '../test/repl-helpers';
+import { expect, useTmpdir, waitEval, waitBus, waitCompletion, fakeTTYProps, readReplLogfile } from '../test/repl-helpers';
 import { CliReplErrors } from './error-codes';
 import { MongoshInternalError } from '@mongosh/errors';
 import http from 'http';
@@ -208,7 +208,8 @@ describe('CliRepl', () => {
     verifyAutocompletion({
       testServer: null,
       wantWatch: true,
-      wantShardDistribution: true
+      wantShardDistribution: true,
+      hasCollectionNames: false
     });
   });
 
@@ -353,7 +354,8 @@ describe('CliRepl', () => {
     verifyAutocompletion({
       testServer: testServer,
       wantWatch: false,
-      wantShardDistribution: false
+      wantShardDistribution: false,
+      hasCollectionNames: true
     });
 
     context('analytics integration', () => {
@@ -439,22 +441,25 @@ describe('CliRepl', () => {
     verifyAutocompletion({
       testServer: startTestServer('not-shared', '--replicaset', '--nodes', '1'),
       wantWatch: true,
-      wantShardDistribution: false
+      wantShardDistribution: false,
+      hasCollectionNames: true
     });
   });
 
   context('with a mongos', () => {
     verifyAutocompletion({
       testServer: startTestServer('not-shared', '--replicaset', '--sharded', '0'),
       wantWatch: true,
-      wantShardDistribution: true
+      wantShardDistribution: true,
+      hasCollectionNames: false // We're only spinning up a mongos here
     });
   });
 
-  function verifyAutocompletion({ testServer, wantWatch, wantShardDistribution }: {
+  function verifyAutocompletion({ testServer, wantWatch, wantShardDistribution, hasCollectionNames }: {
     testServer: MongodSetup | null,
     wantWatch: boolean,
-    wantShardDistribution: boolean
+    wantShardDistribution: boolean,
+    hasCollectionNames: boolean
   }): void {
     describe('autocompletion', () => {
       let cliRepl: CliRepl;
@@ -475,7 +480,7 @@ describe('CliRepl', () => {
       it(`${wantWatch ? 'completes' : 'does not complete'} the watch method`, async() => {
         output = '';
         input.write('db.wat\u0009\u0009');
-        await tick();
+        await waitCompletion(cliRepl.bus);
         if (wantWatch) {
           expect(output).to.include('db.watch');
         } else {
@@ -486,13 +491,28 @@ describe('CliRepl', () => {
       it(`${wantShardDistribution ? 'completes' : 'does not complete'} the getShardDistribution method`, async() => {
         output = '';
         input.write('db.coll.getShardDis\u0009\u0009');
-        await tick();
+        await waitCompletion(cliRepl.bus);
         if (wantShardDistribution) {
           expect(output).to.include('db.coll.getShardDistribution');
         } else {
           expect(output).not.to.include('db.coll.getShardDistribution');
         }
       });
+
+      it('includes collection names', async() => {
+        if (!hasCollectionNames) return;
+        const collname = `testcollection${Date.now()}${(Math.random() * 1000) | 0}`;
+        input.write(`db.${collname}.insertOne({});\n`);
+        await waitEval(cliRepl.bus);
+
+        output = '';
+        input.write('db.testcoll\u0009\u0009');
+        await waitCompletion(cliRepl.bus);
+        expect(output).to.include(collname);
+
+        input.write(`db.${collname}.drop()\n`);
+        await waitEval(cliRepl.bus);
+      });
     });
   }
 });

packages/cli-repl/src/mongosh-repl.ts

@@ -127,6 +127,7 @@ class MongoshNodeRepl implements EvaluationListener {
           (async() => await origReplCompleter(text) || [[]])(),
           (async() => await mongoshCompleter(text))()
         ]);
+        this.bus.emit('mongosh:autocompletion-complete'); // For testing.
         // Remove duplicates, because shell API methods might otherwise show
         // up in both completions.
         const deduped = [...new Set([...replResults, ...mongoshResults])];

packages/cli-repl/test/repl-helpers.ts

@@ -53,6 +53,11 @@ async function waitEval(bus: MongoshBus) {
   await tick();
 }
 
+async function waitCompletion(bus: MongoshBus) {
+  await waitBus(bus, 'mongosh:autocompletion-complete');
+  await tick();
+}
+
 const fakeTTYProps = {
   isTTY: true,
   isRaw: true,
@@ -74,6 +79,7 @@ export {
   tick,
   waitBus,
   waitEval,
+  waitCompletion,
   fakeTTYProps,
   readReplLogfile
 };

packages/shell-api/src/database.ts

@@ -53,6 +53,7 @@ export default class Database extends ShellApiClass {
   _collections: Record<string, Collection>;
   _baseOptions: CommandOperationOptions;
   _session: Session | undefined;
+  _cachedCollectionNames: string[] = [];
 
   constructor(mongo: Mongo, name: string, session?: Session) {
     super();
@@ -138,6 +139,30 @@ export default class Database extends ShellApiClass {
     );
   }
 
+  async _getCollectionNames(): Promise<string[]> {
+    const infos = await this._listCollections({}, { nameOnly: true });
+    this._cachedCollectionNames = infos.map((collection: any) => collection.name);
+    return this._cachedCollectionNames;
+  }
+
+  async _getCollectionNamesForCompletion(): Promise<string[]> {
+    return await Promise.race([
+      (async() => {
+        return await this._getCollectionNames();
+      })(),
+      (async() => {
+        // 200ms should be a good compromise between giving the server a chance
+        // to reply and responsiveness for human perception. It's not the end
+        // of the world if we end up using the cached results; usually, they
+        // are not going to differ from fresh ones, and even if they do, a
+        // subsequent autocompletion request will almost certainly have at least
+        // the new cached results.
+        await new Promise(resolve => setTimeout(resolve, 200).unref());
+        return this._cachedCollectionNames;
+      })()
+    ]);
+  }
+
   async _getLastErrorObj(w?: number|string, wTimeout?: number, j?: boolean): Promise<Document> {
     const cmd = { getlasterror: 1 } as any;
     if (w) {
@@ -177,8 +202,7 @@ export default class Database extends ShellApiClass {
   @returnsPromise
   async getCollectionNames(): Promise<string[]> {
     this._emitDatabaseApiCall('getCollectionNames');
-    const infos = await this._listCollections({}, { nameOnly: true });
-    return infos.map((collection: any) => collection.name);
+    return this._getCollectionNames();
   }
 
   /**

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

@@ -1708,12 +1708,22 @@ describe('Shell API (integration)', function() {
     });
 
     describe('getAutocompleteParameters', () => {
+      beforeEach(async() => {
+        // Make sure the collection is present so it is included in autocompletion.
+        await collection.insertOne({});
+        // Make sure 'database' is the current db in the eyes of the internal state object.
+        internalState.setDbFunc(database);
+      });
+
       it('returns information that is meaningful for autocompletion', async() => {
         const params = await internalState.getAutocompleteParameters();
         expect(params.topology()).to.equal(Topologies.Standalone);
         expect(params.connectionInfo().uri).to.equal(await testServer.connectionString());
         expect(params.connectionInfo().is_atlas).to.equal(false);
         expect(params.connectionInfo().is_localhost).to.equal(true);
+        expect(await database._getCollectionNames()).to.deep.equal(['docs']);
+        expect(await params.getCollectionCompletionsForCurrentDb('d')).to.deep.equal(['docs']);
+        expect(await params.getCollectionCompletionsForCurrentDb('e')).to.deep.equal([]);
       });
     });
   });

packages/shell-api/src/shell-internal-state.spec.ts

@@ -46,8 +46,11 @@ describe('ShellInternalState', () => {
       expect(() => run('db = 42')).to.throw("[COMMON-10002] Cannot reassign 'db' to non-Database type");
     });
 
-    it('allows setting db to a db', async() => {
+    it('allows setting db to a db and causes prefetching', async() => {
+      serviceProvider.listCollections
+        .resolves([ { name: 'coll1' }, { name: 'coll2' } ]);
       expect(run('db = db.getSiblingDB("moo"); db.getName()')).to.equal('moo');
+      expect(serviceProvider.listCollections.calledWith('moo', {}, { nameOnly: true })).to.equal(true);
     });
   });
 });