feat: Suggest quoted entities for MongoDB (#280)

* feat: support collections suggestion * feat: support user suggestions * refactor: rename variables * feat: suggest collections for indexInformation method
gravity-ui · Jan 27, 2025 · 4cc7db0 · 4cc7db0
1 parent 71acf78
commit 4cc7db0
Show file tree

Hide file tree

Showing 12 changed files with 1,598 additions and 1,443 deletions.
diff --git a/src/autocomplete/databases/mongo/generated/MongoParser.interp b/src/autocomplete/databases/mongo/generated/MongoParser.interp
diff --git a/src/autocomplete/databases/mongo/generated/MongoParser.ts b/src/autocomplete/databases/mongo/generated/MongoParser.ts
diff --git a/src/autocomplete/databases/mongo/generated/MongoParserVisitor.ts b/src/autocomplete/databases/mongo/generated/MongoParserVisitor.ts
@@ -49,6 +49,7 @@ import { DatabaseIndexInformationArgument1Context } from "./MongoParser.js";
 import { DatabaseIndexInformationArgument2Context } from "./MongoParser.js";
 import { DatabaseRemoveUserMethodContext } from "./MongoParser.js";
 import { DatabaseRemoveUserArgument1Context } from "./MongoParser.js";
+import { QuotedUsernameContext } from "./MongoParser.js";
 import { DatabaseRemoveUserArgument2Context } from "./MongoParser.js";
 import { DatabaseCreateIndexMethodContext } from "./MongoParser.js";
 import { DatabaseCreateIndexArgument3Context } from "./MongoParser.js";
@@ -425,6 +426,12 @@ export class MongoParserVisitor<Result> extends AbstractParseTreeVisitor<Result>
      * @return the visitor result
      */
     visitDatabaseRemoveUserArgument1?: (ctx: DatabaseRemoveUserArgument1Context) => Result;
+    /**
+     * Visit a parse tree produced by `MongoParser.quotedUsername`.
+     * @param ctx the parse tree
+     * @return the visitor result
+     */
+    visitQuotedUsername?: (ctx: QuotedUsernameContext) => Result;
     /**
      * Visit a parse tree produced by `MongoParser.databaseRemoveUserArgument2`.
      * @param ctx the parse tree

diff --git a/src/autocomplete/databases/mongo/grammar/MongoParser.g4 b/src/autocomplete/databases/mongo/grammar/MongoParser.g4
@@ -171,7 +171,7 @@ databaseIndexInformationMethod
     ;
 
 databaseIndexInformationArgument1
-    : STRING
+    : quotedCollectionName
     ;
 
 databaseIndexInformationArgument2
@@ -183,6 +183,10 @@ databaseRemoveUserMethod
     ;
 
 databaseRemoveUserArgument1
+    : quotedUsername
+    ;
+
+quotedUsername
     : STRING
     ;
 
@@ -225,7 +229,7 @@ databaseDropCollectionMethod
     ;
 
 databaseDropCollectionArgument1
-    : STRING
+    : quotedCollectionName
     ;
 
 databaseDropCollectionArgument2
@@ -239,7 +243,7 @@ databaseRenameCollectionMethod
     ;
 
 databaseRenameCollectionArgument1
-    : STRING
+    : quotedCollectionName
     ;
 
 databaseRenameCollectionArgument2

diff --git a/src/autocomplete/databases/mongo/index.ts b/src/autocomplete/databases/mongo/index.ts
@@ -11,7 +11,9 @@ import {
 export * from './mongo-extract-commands';
 
 export interface MongoAutocompleteResult extends SqlAutocompleteResult {
+    suggestQuotedCollections?: boolean;
     suggestCollections?: boolean;
+    suggestQuotedUsers?: boolean;
 }
 
 export function parseMongoQueryWithoutCursor(

diff --git a/src/autocomplete/databases/mongo/mongo-autocomplete.ts b/src/autocomplete/databases/mongo/mongo-autocomplete.ts
@@ -32,13 +32,19 @@ function getIgnoredTokens(): number[] {
 
 const ignoredTokens = new Set(getIgnoredTokens());
 
-const rulesToVisit = new Set([MongoParser.RULE_collectionName]);
+const rulesToVisit = new Set([
+    MongoParser.RULE_collectionName,
+    MongoParser.RULE_quotedCollectionName,
+    MongoParser.RULE_quotedUsername,
+]);
 
 function processVisitedRules(
     rules: c3.CandidatesCollection['rules'],
     cursorTokenIndex: number,
 ): ProcessVisitedRulesResult<MongoAutocompleteResult> {
+    let suggestQuotedCollections;
     let suggestCollections;
+    let suggestQuotedUsers;
 
     for (const [ruleId, rule] of rules) {
         if (!isStartingToWriteRule(cursorTokenIndex, rule)) {
@@ -50,10 +56,22 @@ function processVisitedRules(
                 suggestCollections = true;
                 break;
             }
+            case MongoParser.RULE_quotedCollectionName: {
+                suggestQuotedCollections = true;
+                break;
+            }
+            case MongoParser.RULE_quotedUsername: {
+                suggestQuotedUsers = true;
+                break;
+            }
         }
     }
 
-    return {suggestCollections};
+    return {
+        suggestQuotedCollections,
+        suggestCollections,
+        suggestQuotedUsers,
+    };
 }
 
 export function getParseTree(parser: MongoParser): ParseTree {

diff --git a/src/autocomplete/databases/mongo/tests/admin/admin.test.ts b/src/autocomplete/databases/mongo/tests/admin/admin.test.ts
@@ -61,6 +61,14 @@ test('should not report errors on extended admin statements', () => {
     expect(autocompleteResult.errors).toHaveLength(0);
 });
 
+test('should suggest users on first removeUser argument', () => {
+    const autocompleteResult = parseMongoQueryWithCursor(`
+        db.admin().removeUser(|
+    `);
+
+    expect(autocompleteResult.suggestQuotedUsers).toEqual(true);
+});
+
 test('should suggest keywords after db.admin().', () => {
     const autocompleteResult = parseMongoQueryWithCursor('db.admin().|');
 

diff --git a/src/autocomplete/databases/mongo/tests/db/db.test.ts b/src/autocomplete/databases/mongo/tests/db/db.test.ts
@@ -14,6 +14,24 @@ test('should not report errors on multiple statements', () => {
     expect(autocompleteResult.errors).toHaveLength(0);
 });
 
+test('should suggest collections without quotes after db', () => {
+    const autocompleteResult = parseMongoQueryWithCursor(`
+        db.|
+    `);
+
+    expect(autocompleteResult.suggestCollections).toEqual(true);
+    expect(autocompleteResult.suggestQuotedCollections).toBeUndefined();
+});
+
+test('should suggest collections in collection method', () => {
+    const autocompleteResult = parseMongoQueryWithCursor(`
+        db.collection(|
+    `);
+
+    expect(autocompleteResult.suggestCollections).toBeUndefined();
+    expect(autocompleteResult.suggestQuotedCollections).toEqual(true);
+});
+
 test('should not report errors on three statements', () => {
     const autocompleteResult = parseMongoQueryWithoutCursor(`
         db.test_collection.find();

diff --git a/src/autocomplete/databases/mongo/tests/drop/drop-collection.test.ts b/src/autocomplete/databases/mongo/tests/drop/drop-collection.test.ts
@@ -1,4 +1,9 @@
-import {Command, extractMongoCommandsFromQuery, parseMongoQueryWithoutCursor} from '../..';
+import {
+    Command,
+    extractMongoCommandsFromQuery,
+    parseMongoQueryWithCursor,
+    parseMongoQueryWithoutCursor,
+} from '../..';
 
 test('should not report errors on dropCollection statement', () => {
     const autocompleteResult = parseMongoQueryWithoutCursor(`
@@ -21,6 +26,14 @@ test('should not report errors on extended dropCollection statement', () => {
     expect(autocompleteResult.errors).toHaveLength(0);
 });
 
+test('should suggest collections on first dropCollection argument', () => {
+    const autocompleteResult = parseMongoQueryWithCursor(`
+        db.dropCollection(|
+    `);
+
+    expect(autocompleteResult.suggestQuotedCollections).toEqual(true);
+});
+
 test('should extract dropCollection commands properly', () => {
     const result = extractMongoCommandsFromQuery(`
         db.dropCollection('test_collection_name');

diff --git a/src/autocomplete/databases/mongo/tests/index/index-information.test.ts b/src/autocomplete/databases/mongo/tests/index/index-information.test.ts
@@ -1,4 +1,9 @@
-import {Command, extractMongoCommandsFromQuery, parseMongoQueryWithoutCursor} from '../..';
+import {
+    Command,
+    extractMongoCommandsFromQuery,
+    parseMongoQueryWithCursor,
+    parseMongoQueryWithoutCursor,
+} from '../..';
 
 test('should not report errors on indexInformation statement', () => {
     const autocompleteResult = parseMongoQueryWithoutCursor(`
@@ -34,6 +39,14 @@ test('should not report errors on extended indexInformation statement', () => {
     expect(autocompleteResult.errors).toHaveLength(0);
 });
 
+test('should suggest quoted collections on first argument of indexInformation method in database context', () => {
+    const autocompleteResult = parseMongoQueryWithCursor(`
+        db.indexInformation(|
+    `);
+
+    expect(autocompleteResult.suggestQuotedCollections).toEqual(true);
+});
+
 test('should extract indexInformation commands properly', () => {
     const result = extractMongoCommandsFromQuery(`
         db.test_collection.indexInformation();

diff --git a/src/autocomplete/databases/mongo/tests/remove/remove-user.test.ts b/src/autocomplete/databases/mongo/tests/remove/remove-user.test.ts
@@ -1,4 +1,9 @@
-import {Command, extractMongoCommandsFromQuery, parseMongoQueryWithoutCursor} from '../..';
+import {
+    Command,
+    extractMongoCommandsFromQuery,
+    parseMongoQueryWithCursor,
+    parseMongoQueryWithoutCursor,
+} from '../..';
 
 test('should not report errors on removeUser statement', () => {
     const autocompleteResult = parseMongoQueryWithoutCursor(`
@@ -21,6 +26,14 @@ test('should not report errors on extended removeUser statement', () => {
     expect(autocompleteResult.errors).toHaveLength(0);
 });
 
+test('should suggest users on first removeUser argument', () => {
+    const autocompleteResult = parseMongoQueryWithCursor(`
+        db.removeUser(|
+    `);
+
+    expect(autocompleteResult.suggestQuotedUsers).toEqual(true);
+});
+
 test('should extract removeUser commands properly', () => {
     const result = extractMongoCommandsFromQuery(`
         db.removeUser('test_user');

diff --git a/src/autocomplete/databases/mongo/tests/rename/rename-collection.test.ts b/src/autocomplete/databases/mongo/tests/rename/rename-collection.test.ts
@@ -1,4 +1,9 @@
-import {Command, extractMongoCommandsFromQuery, parseMongoQueryWithoutCursor} from '../..';
+import {
+    Command,
+    extractMongoCommandsFromQuery,
+    parseMongoQueryWithCursor,
+    parseMongoQueryWithoutCursor,
+} from '../..';
 
 test('should not report errors on renameCollection statement', () => {
     const autocompleteResult = parseMongoQueryWithoutCursor(`
@@ -8,6 +13,14 @@ test('should not report errors on renameCollection statement', () => {
     expect(autocompleteResult.errors).toHaveLength(0);
 });
 
+test('should suggest collections on first renameCollection argument', () => {
+    const autocompleteResult = parseMongoQueryWithCursor(`
+        db.renameCollection(|
+    `);
+
+    expect(autocompleteResult.suggestQuotedCollections).toEqual(true);
+});
+
 test('should not report errors on extended renameCollection statement', () => {
     const autocompleteResult = parseMongoQueryWithoutCursor(`
         db.renameCollection(