fix(autocomplete): fixes autocomplete not working for db.aggregate

packages/autocomplete/package.json

@@ -31,7 +31,7 @@
     "mocha": "^7.1.2"
   },
   "dependencies": {
-    "@mongodb-js/mongodb-constants": "^0.2.1",
+    "@mongodb-js/mongodb-constants": "^0.2.2",
     "@mongosh/shell-api": "0.0.0-dev.0",
     "semver": "^7.3.2"
   }

packages/autocomplete/src/index.spec.ts

@@ -5,6 +5,17 @@ import { expect } from 'chai';
 
 let collections: string[];
 let databases: string[];
+const standalone600 = {
+  topology: () => Topologies.Standalone,
+  apiVersionInfo: () => undefined,
+  connectionInfo: () => ({
+    is_atlas: false,
+    is_data_federation: false,
+    server_version: '6.0.0'
+  }),
+  getCollectionCompletionsForCurrentDb: () => collections,
+  getDatabaseCompletions: () => databases
+};
 const standalone440 = {
   topology: () => Topologies.Standalone,
   apiVersionInfo: () => undefined,
@@ -324,6 +335,53 @@ describe('completer.completer', () => {
     });
   });
 
+  context('when context is db aggregation query', () => {
+    it('has several matches for db level stages', async() => {
+      const query = 'db.aggregate([{';
+      expect(await completer(standalone440, query)).to.deep.equal([
+        [
+          'db.aggregate([{$changeStream',
+          'db.aggregate([{$currentOp',
+          'db.aggregate([{$listLocalSessions',
+        ],
+        query
+      ]);
+      expect(await completer(standalone600, query)).to.deep.equal([
+        [
+          'db.aggregate([{$documents',
+          'db.aggregate([{$changeStream',
+          'db.aggregate([{$currentOp',
+          'db.aggregate([{$listLocalSessions',
+        ],
+        query
+      ]);
+    });
+
+    it('does not have a match', async() => {
+      const query = 'db.aggregate([{$mat';
+      expect(await completer(standalone440, query)).to.deep.equal([[], query]);
+    });
+
+    it('matches a db aggregation stage', async() => {
+      const query = 'db.aggregate([{$lis';
+      expect(await completer(standalone440, query)).to.deep.equal([
+        ['db.aggregate([{$listLocalSessions'],
+        query
+      ]);
+    });
+
+    it('completes the followup stages', async() => {
+      const query = 'db.aggregate([{$currentOp: {}}, {$ma';
+      expect(await completer(standalone440, query)).to.deep.equal([
+        [
+          'db.aggregate([{$currentOp: {}}, {$map',
+          'db.aggregate([{$currentOp: {}}, {$match'
+        ],
+        query
+      ]);
+    });
+  });
+
   context('when context is aggregation query', () => {
     it('has several matches', async() => {
       const i = 'db.shipwrecks.aggregate([ { $so';

packages/autocomplete/src/index.ts

@@ -11,7 +11,8 @@ import {
   BSON_TYPES,
   ATLAS,
   ADL,
-  ON_PREM
+  ON_PREM,
+  DATABASE
 } from '@mongodb-js/mongodb-constants';
 
 type TypeSignatureAttributes = { [key: string]: TypeSignature };
@@ -49,6 +50,12 @@ export const MATCH_COMPLETIONS = ([] as AnyCompletions).concat(
   BSON_TYPES
 );
 
+// Note: The following list is not a list of all the completions
+// for `db.aggregate` but only for the first stage of `db.aggregate`.
+const DB_AGGREGATE_COMPLETIONS = STAGE_OPERATORS.filter(({ namespaces }) => {
+  return namespaces.length === 1 && namespaces[0] === DATABASE;
+});
+
 /**
  * The project stage operator.
  */
@@ -116,6 +123,23 @@ async function completer(params: AutocompleteParameters, line: string): Promise<
     const hits = filterShellAPI(params, SHELL_COMPLETIONS, elToComplete);
     return [hits.length ? hits : [], line];
   } else if (firstLineEl.match(/\bdb\b/) && splitLine.length === 2) {
+    if (elToComplete.match(/aggregate\s*\(\s*\[\s*\{\s*/)) {
+      const splitQuery = line.split('{');
+      const prefix = splitQuery.pop()?.trim() || '';
+      const command: string = prefix ? line.split(prefix).shift() as string : line;
+      const suggestFirstStage = splitQuery.length <= 2;
+
+      const expressions = suggestFirstStage
+        // First stage in `db.aggregate` form can only be 'db' namespaced stages
+        ? DB_AGGREGATE_COMPLETIONS
+        : [
+          ...BASE_COMPLETIONS,
+          ...getStageAccumulators(params, elToComplete)
+        ];
+
+      const hits = filterQueries(params, expressions, prefix, command);
+      return [hits.length ? hits : [], line];
+    }
     // We're seeing something like 'db.foo' and expand that to all methods on
     // db which start with 'foo' and all collections on the current db that
     // start with 'foo'.
@@ -154,10 +178,10 @@ async function completer(params: AutocompleteParameters, line: string): Promise<
       let expressions;
       if (splitLine[2].match(/\baggregate\b/)) {
         // aggregation needs extra accumulators to autocomplete properly
-        expressions = ([] as AnyCompletions).concat(
-          BASE_COMPLETIONS,
-          getStageAccumulators(params, elToComplete)
-        );
+        expressions = [
+          ...BASE_COMPLETIONS,
+          ...getStageAccumulators(params, elToComplete)
+        ];
       } else {
         // collection querying just needs MATCH COMPLETIONS
         expressions = MATCH_COMPLETIONS;