feat: system variables and fields completion VSCODE-377, VSCODE-393 (#…

…500) * feat: system variables and fields completion VSCODE-377, VSCODE-393 * test: wrap aggregation input into array
mongodb-js · Mar 27, 2023 · 6dfc573 · 6dfc573
1 parent d50581b
commit 6dfc573
Show file tree

Hide file tree

Showing 5 changed files with 96 additions and 32 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -968,7 +968,7 @@
     "@iconify/react": "^1.1.4",
     "@leafygreen-ui/logo": "^6.3.0",
     "@leafygreen-ui/toggle": "^7.0.5",
-    "@mongodb-js/mongodb-constants": "^0.3.0",
+    "@mongodb-js/mongodb-constants": "^0.4.0",
     "@mongosh/browser-runtime-electron": "^1.8.0",
     "@mongosh/i18n": "^1.8.0",
     "@mongosh/service-provider-server": "^1.8.0",
@@ -1057,7 +1057,7 @@
     "ora": "^5.4.1",
     "postcss-loader": "^3.0.0",
     "pre-commit": "^1.2.2",
-    "prettier": "^2.8.6",
+    "prettier": "^2.8.7",
     "process": "^0.11.10",
     "semver": "^7.3.8",
     "sinon": "^9.2.4",

diff --git a/src/language/mongoDBService.ts b/src/language/mongoDBService.ts
@@ -587,7 +587,7 @@ export default class MongoDBService {
   /**
    * If the current node is 'db.collection.find({ _id: <trigger>});'.
    */
-  _provideBSONCompletionItems(state: CompletionState) {
+  _provideIdentifierObjectValueCompletionItems(state: CompletionState) {
     if (state.isIdentifierObjectValue) {
       this._connection.console.log('VISITOR found bson completions');
       return getFilteredCompletions({ meta: ['bson'] }).map((item) => {
@@ -611,24 +611,28 @@ export default class MongoDBService {
   }
 
   /**
-   * If the current node is 'db.collection.find({ $expr: { $gt: [{ $getField: { $literal: '<trigger>' } }, 200] } });'.
+   * If the current node is 'db.collection.find({ field: "<trigger>" });'.
    */
-  _provideFieldReferenceCompletionItems(state: CompletionState) {
+  _provideTextObjectValueCompletionItems(state: CompletionState) {
     if (state.isTextObjectValue) {
       const fields =
         this._fields[`${state.databaseName}.${state.collectionName}`];
       this._connection.console.log('VISITOR found field reference completions');
 
-      return getFilteredCompletions({ fields, meta: ['field:reference'] }).map(
-        (item) => {
-          return {
-            label: item.value,
-            kind: CompletionItemKind.Reference,
-            preselect: true,
-            detail: item.description,
-          };
-        }
-      );
+      return getFilteredCompletions({
+        fields,
+        meta: ['field:reference', 'variable:*'],
+      }).map((item) => {
+        return {
+          label: item.value,
+          kind:
+            item.meta === 'field:reference'
+              ? CompletionItemKind.Reference
+              : CompletionItemKind.Variable,
+          preselect: true,
+          detail: item.description,
+        };
+      });
     }
   }
 
@@ -750,8 +754,8 @@ export default class MongoDBService {
       this._provideStageCompletionItems.bind(this, state),
       this._provideQueryOperatorCompletionItems.bind(this, state),
       this._provideAggregationOperatorCompletionItems.bind(this, state),
-      this._provideBSONCompletionItems.bind(this, state),
-      this._provideFieldReferenceCompletionItems.bind(this, state),
+      this._provideIdentifierObjectValueCompletionItems.bind(this, state),
+      this._provideTextObjectValueCompletionItems.bind(this, state),
       this._provideCollectionSymbolCompletionItems.bind(this, state),
       this._provideFindCursorCompletionItems.bind(this, state),
       this._provideAggregationCursorCompletionItems.bind(this, state),

diff --git a/src/language/visitor.ts b/src/language/visitor.ts
@@ -561,7 +561,9 @@ export class Visitor {
     }
   }
 
-  _checkHasCollectionName(node: babel.types.MemberExpression): void {
+  _checkHasCollectionNameMemberExpression(
+    node: babel.types.MemberExpression
+  ): void {
     if (
       node.object.type === 'MemberExpression' &&
       node.object.object.type === 'Identifier' &&
@@ -575,6 +577,26 @@ export class Visitor {
     }
   }
 
+  _checkHasCollectionNameCallExpression(node: babel.types.MemberExpression) {
+    if (
+      node.object.type === 'CallExpression' &&
+      node.object.callee.type === 'MemberExpression' &&
+      node.object.callee.object.type === 'Identifier' &&
+      node.object.callee.object.name === 'db' &&
+      node.object.callee.property.type === 'Identifier' &&
+      node.object.callee.property.name === 'getCollection' &&
+      node.object.arguments.length === 1 &&
+      node.object.arguments[0].type === 'StringLiteral'
+    ) {
+      this._state.collectionName = node.object.arguments[0].value;
+    }
+  }
+
+  _checkHasCollectionName(node: babel.types.MemberExpression): void {
+    this._checkHasCollectionNameMemberExpression(node);
+    this._checkHasCollectionNameCallExpression(node);
+  }
+
   _checkIsCollectionMemberExpression(node: babel.types.MemberExpression): void {
     if (
       node.object.type === 'MemberExpression' &&

diff --git a/src/test/suite/language/mongoDBService.test.ts b/src/test/suite/language/mongoDBService.test.ts
@@ -590,6 +590,30 @@ suite('MongoDBService Test Suite', () => {
       expect(completion).to.have.property('detail');
     });
 
+    test('provide system variable completion in find', async () => {
+      const result = await testMongoDBService.provideCompletionItems(
+        'db.collection.find({ _id: "$$N" });',
+        { line: 0, character: 30 }
+      );
+      const completion = result.find(
+        (item: CompletionItem) => item.label === '$$NOW'
+      );
+
+      expect(completion).to.have.property('kind', CompletionItemKind.Variable);
+    });
+
+    test('provide system variable completion in aggregate', async () => {
+      const result = await testMongoDBService.provideCompletionItems(
+        'db.collection.aggregate([{ $match: { _id: "$$R" }}]);',
+        { line: 0, character: 46 }
+      );
+      const completion = result.find(
+        (item: CompletionItem) => item.label === '$$ROOT'
+      );
+
+      expect(completion).to.have.property('kind', CompletionItemKind.Variable);
+    });
+
     test('provide field reference completion in find when has db', async () => {
       testMongoDBService._cacheFields('test.collection', ['price']);
 
@@ -630,6 +654,20 @@ suite('MongoDBService Test Suite', () => {
       expect(completion).to.have.property('kind', CompletionItemKind.Reference);
     });
 
+    test('provide field reference completion in aggregate when collection is specified via getCollection', async () => {
+      testMongoDBService._cacheFields('test.collection', ['price']);
+
+      const result = await testMongoDBService.provideCompletionItems(
+        "use('test'); db.getCollection('collection').aggregate([{ $match: '$p' }]);",
+        { line: 0, character: 68 }
+      );
+      const completion = result.find(
+        (item: CompletionItem) => item.label === '$price'
+      );
+
+      expect(completion).to.have.property('kind', CompletionItemKind.Reference);
+    });
+
     test('provide aggregation expression completion for other than $match stages', async () => {
       const result = await testMongoDBService.provideCompletionItems(
         'db.collection.aggregate([{ $project: { yearMonthDayUTC: { $d } } }]);',