feat: mark as entityCollecting in getAllEntities context to allow empty column

Author: JackWang032

src/parser/common/basicSQL.ts

@@ -134,6 +134,7 @@ export abstract class BasicSQL<
     public createParser(input: string, errorListener?: ErrorListener) {
         const lexer = this.createLexer(input, errorListener);
         const tokenStream = new CommonTokenStream(lexer);
+        tokenStream.fill();
         const parser = this.createParserFromTokenStream(tokenStream);
         parser.interpreter.predictionMode = PredictionMode.SLL;
         if (errorListener) {
@@ -557,29 +558,28 @@ export abstract class BasicSQL<
     }
 
     public getAllEntities(input: string, caretPosition?: CaretPosition): EntityContext[] | null {
-        const allTokens = this.getAllTokens(input);
+        /**
+         * Create a new parser to generate brand new parse tree.
+         * And the new parse tree should not effect cached parse tree which is used by validate and getSuggestionAtCaretPosition method.
+         */
+        const parser = this.createParser(input);
+        const allTokens = (parser.tokenStream as CommonTokenStream).getTokens();
+
         const caretTokenIndex = caretPosition
             ? findCaretTokenIndex(caretPosition, allTokens)
             : void 0;
-
         const collectListener = this.createEntityCollector(input, allTokens, caretTokenIndex);
-        // const parser = this.createParserWithCache(input);
 
-        // parser.entityCollecting = true;
-        // if(caretPosition) {
-        //     const allTokens = this.getAllTokens(input);
-        //     const tokenIndex = findCaretTokenIndex(caretPosition, allTokens);
-        //     parser.caretTokenIndex = tokenIndex;
-        // }
-
-        // const parseTree = parser.program();
-
-        const parseTree = this.parseWithCache(input);
+        parser.entityCollecting = true;
+        if (caretPosition && caretTokenIndex !== undefined) {
+            parser.caretTokenIndex = caretTokenIndex;
+        }
 
+        const parseTree = parser.program();
         this.listen(collectListener, parseTree);
 
-        // parser.caretTokenIndex = -1;
-        // parser.entityCollecting = false;
+        parser.caretTokenIndex = -1;
+        parser.entityCollecting = false;
 
         return collectListener.getEntities();
     }

test/parser/impala/suggestion/suggestionWithEntity.test.ts

@@ -70,15 +70,14 @@ describe('Impala SQL Syntax Suggestion with collect entity', () => {
         expect(suggestion?.wordRanges.map((token) => token.text)).toEqual([]);
 
         const entities = impala.getAllEntities(sql, pos);
-        expect(entities.length).toBe(1);
+        expect(entities.length).toBe(2);
         expect(entities[0].text).toBe('insert_tb');
         expect(entities[0].entityContextType).toBe(EntityContextType.TABLE);
         expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
 
-        // TODO:
-        // expect(entities[1].text).toBe('from_tb');
-        // expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
-        // expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
+        expect(entities[1].text).toBe('from_tb');
+        expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
+        expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
     });
 
     test('insert into table as select with trailing comma', () => {
@@ -121,14 +120,14 @@ describe('Impala SQL Syntax Suggestion with collect entity', () => {
         expect(suggestion?.wordRanges.map((token) => token.text)).toEqual([]);
 
         const entities = impala.getAllEntities(sql, pos);
-        expect(entities.length).toBe(1);
+        expect(entities.length).toBe(2);
         expect(entities[0].text).toBe('sorted_census_data');
         expect(entities[0].entityContextType).toBe(EntityContextType.TABLE_CREATE);
         expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
-        // TODO:
-        // expect(entities[1].text).toBe('unsorted_census_data');
-        // expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
-        // expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
+
+        expect(entities[1].text).toBe('unsorted_census_data');
+        expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
+        expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
     });
 
     test('create table as select with trailing comma', () => {

test/parser/trino/suggestion/suggestionWithEntity.test.ts

@@ -70,15 +70,14 @@ describe('PostgreSql Syntax Suggestion with collect entity', () => {
         expect(suggestion?.wordRanges.map((token) => token.text)).toEqual([]);
 
         const entities = trino.getAllEntities(sql, pos);
-        expect(entities.length).toBe(1);
+        expect(entities.length).toBe(2);
         expect(entities[0].text).toBe('insert_tb');
         expect(entities[0].entityContextType).toBe(EntityContextType.TABLE);
         expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
 
-        // TODO:
-        // expect(entities[1].text).toBe('from_tb');
-        // expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
-        // expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
+        expect(entities[1].text).toBe('from_tb');
+        expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
+        expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
     });
 
     test('insert into table as select with trailing comma', () => {
@@ -121,15 +120,14 @@ describe('PostgreSql Syntax Suggestion with collect entity', () => {
         expect(suggestion?.wordRanges.map((token) => token.text)).toEqual([]);
 
         const entities = trino.getAllEntities(sql, pos);
-        expect(entities.length).toBe(1);
+        expect(entities.length).toBe(2);
         expect(entities[0].text).toBe('sorted_census_data');
         expect(entities[0].entityContextType).toBe(EntityContextType.TABLE_CREATE);
         expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
 
-        // TODO:
-        // expect(entities[1].text).toBe('unsorted_census_data');
-        // expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
-        // expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
+        expect(entities[1].text).toBe('unsorted_census_data');
+        expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
+        expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
     });
 
     test('create table as select with trailing comma', () => {