fix: update isContainCaret judgment when caret position token is whit…

src/parser/common/basicSQL.ts

@@ -83,6 +83,7 @@ export abstract class BasicSQL<
      */
     protected abstract createEntityCollector(
         input: string,
+        allTokens?: Token[],
         caretTokenIndex?: number
     ): EntityCollector;
 
@@ -378,7 +379,7 @@ export abstract class BasicSQL<
             ? findCaretTokenIndex(caretPosition, allTokens)
             : void 0;
 
-        const collectListener = this.createEntityCollector(input, caretTokenIndex);
+        const collectListener = this.createEntityCollector(input, allTokens, caretTokenIndex);
         // const parser = this.createParserWithCache(input);
 
         // parser.entityCollecting = true;

src/parser/common/entityCollector.ts

@@ -1,4 +1,4 @@
-import { ParserRuleContext } from 'antlr4ng';
+import { ParserRuleContext, Token } from 'antlr4ng';
 import { EntityContextType } from './types';
 import { WordPosition, TextPosition } from './textAndWord';
 import { ctxToText, ctxToWord } from './textAndWord';
@@ -96,15 +96,17 @@ export function toEntityContext(
  * @todo: [may be need] Combine the entities in each clause.
  */
 export abstract class EntityCollector {
-    constructor(input: string, caretTokenIndex?: number) {
+    constructor(input: string, allTokens?: Token[], caretTokenIndex?: number) {
         this._input = input;
+        this._allTokens = allTokens || [];
         this._caretTokenIndex = caretTokenIndex ?? -1;
         this._entitiesSet = new Set();
         this._stmtStack = new SimpleStack();
         this._entityStack = new SimpleStack();
         this._rootStmt = null;
     }
     private readonly _input: string;
+    private readonly _allTokens: Token[];
     private readonly _caretTokenIndex: number;
     private readonly _entitiesSet: Set<EntityContext>;
     /** Staging statements that have already entered. */
@@ -136,14 +138,31 @@ export abstract class EntityCollector {
         this._rootStmt = null;
     }
 
+    /**
+     * The antlr4 will ignore hidden tokens, if we type whitespace at the end of a statement,
+     * the whitespace token will not as stop token, so we consider the whitespace token as a part of the nonhidden token in front of it
+     */
+    protected getPrevNonHiddenTokenIndex(caretTokenIndex: number) {
+        if (this._allTokens[caretTokenIndex].channel !== Token.HIDDEN_CHANNEL)
+            return caretTokenIndex;
+        for (let i = caretTokenIndex - 1; i >= 0; i--) {
+            const token = this._allTokens[i];
+            if (token.channel !== Token.HIDDEN_CHANNEL) {
+                // If prev nonhidden token is ';', the current token does not belong to any statement.
+                return token.text === ';' ? +Infinity : token.tokenIndex;
+            }
+        }
+        return +Infinity;
+    }
+
     protected pushStmt(ctx: ParserRuleContext, type: StmtContextType) {
         let isContainCaret: boolean | undefined;
         if (this._caretTokenIndex >= 0) {
             isContainCaret =
                 !!ctx.start &&
                 !!ctx.stop &&
                 ctx.start.tokenIndex <= this._caretTokenIndex &&
-                ctx.stop.tokenIndex >= this._caretTokenIndex;
+                ctx.stop.tokenIndex >= this.getPrevNonHiddenTokenIndex(this._caretTokenIndex);
         }
         const stmtContext = toStmtContext(
             ctx,

src/parser/flink/index.ts

@@ -37,8 +37,8 @@ export class FlinkSQL extends BasicSQL<FlinkSqlLexer, ProgramContext, FlinkSqlPa
         return new FlinkSqlSplitListener();
     }
 
-    protected createEntityCollector(input: string, caretTokenIndex?: number) {
-        return new FlinkEntityCollector(input, caretTokenIndex);
+    protected createEntityCollector(input: string, allTokens?: Token[], caretTokenIndex?: number) {
+        return new FlinkEntityCollector(input, allTokens, caretTokenIndex);
     }
 
     protected processCandidates(

src/parser/hive/index.ts

@@ -38,8 +38,8 @@ export class HiveSQL extends BasicSQL<HiveSqlLexer, ProgramContext, HiveSqlParse
         return new HiveSqlSplitListener();
     }
 
-    protected createEntityCollector(input: string, caretTokenIndex?: number) {
-        return new HiveEntityCollector(input, caretTokenIndex);
+    protected createEntityCollector(input: string, allTokens?: Token[], caretTokenIndex?: number) {
+        return new HiveEntityCollector(input, allTokens, caretTokenIndex);
     }
 
     protected processCandidates(

src/parser/impala/index.ts

@@ -36,8 +36,8 @@ export class ImpalaSQL extends BasicSQL<ImpalaSqlLexer, ProgramContext, ImpalaSq
         return new ImpalaSqlSplitListener();
     }
 
-    protected createEntityCollector(input: string, caretTokenIndex?: number) {
-        return new ImpalaEntityCollector(input, caretTokenIndex);
+    protected createEntityCollector(input: string, allTokens?: Token[], caretTokenIndex?: number) {
+        return new ImpalaEntityCollector(input, allTokens, caretTokenIndex);
     }
 
     protected processCandidates(

src/parser/mysql/index.ts

@@ -36,8 +36,8 @@ export class MySQL extends BasicSQL<MySqlLexer, ProgramContext, MySqlParser> {
         return new MysqlSplitListener();
     }
 
-    protected createEntityCollector(input: string, caretTokenIndex?: number) {
-        return new MySqlEntityCollector(input, caretTokenIndex);
+    protected createEntityCollector(input: string, allTokens?: Token[], caretTokenIndex?: number) {
+        return new MySqlEntityCollector(input, allTokens, caretTokenIndex);
     }
 
     protected processCandidates(

src/parser/postgresql/index.ts

@@ -41,8 +41,8 @@ export class PostgreSQL extends BasicSQL<PostgreSqlLexer, ProgramContext, Postgr
         return new PostgreSqlSplitListener();
     }
 
-    protected createEntityCollector(input: string, caretTokenIndex?: number) {
-        return new PostgreSqlEntityCollector(input, caretTokenIndex);
+    protected createEntityCollector(input: string, allTokens?: Token[], caretTokenIndex?: number) {
+        return new PostgreSqlEntityCollector(input, allTokens, caretTokenIndex);
     }
 
     protected processCandidates(

src/parser/spark/index.ts

@@ -36,8 +36,8 @@ export class SparkSQL extends BasicSQL<SparkSqlLexer, ProgramContext, SparkSqlPa
         return new SparkSqlSplitListener();
     }
 
-    protected createEntityCollector(input: string, caretTokenIndex?: number) {
-        return new SparkEntityCollector(input, caretTokenIndex);
+    protected createEntityCollector(input: string, allTokens?: Token[], caretTokenIndex?: number) {
+        return new SparkEntityCollector(input, allTokens, caretTokenIndex);
     }
 
     protected processCandidates(

src/parser/trino/index.ts

@@ -23,8 +23,8 @@ export class TrinoSQL extends BasicSQL<TrinoSqlLexer, ProgramContext, TrinoSqlPa
         return new TrinoSqlSplitListener();
     }
 
-    protected createEntityCollector(input: string, caretTokenIndex?: number) {
-        return new TrinoEntityCollector(input, caretTokenIndex);
+    protected createEntityCollector(input: string, allTokens?: Token[], caretTokenIndex?: number) {
+        return new TrinoEntityCollector(input, allTokens, caretTokenIndex);
     }
 
     protected preferredRules: Set<number> = new Set([

test/parser/flink/suggestion/fixtures/suggestionWithEntity.sql

@@ -8,4 +8,8 @@ INSERT INTO insert_tb PARTITION (country, state) SELECT col1, col2, country, sta
 
 CREATE TABLE IF NOT EXISTS derived_table WITH ('connector' = 'kafka')  AS SELECT  FROM origin_table;
 
-CREATE TABLE IF NOT EXISTS derived_table WITH ('connector' = 'kafka')  AS SELECT id,  FROM origin_table;
+CREATE TABLE IF NOT EXISTS derived_table WITH ('connector' = 'kafka')  AS SELECT id,  FROM origin_table;
+
+SELECT id FROM tb WHERE 
+
+SELECT id FROM tb GROUP BY  ;  

test/parser/flink/suggestion/suggestionWithEntity.test.ts

@@ -157,4 +157,24 @@ describe('Flink SQL Syntax Suggestion with collect entity', () => {
         expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
         expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
     });
+
+    test('isContainCaret should be truthy if caret position is whitespace at the end of statement', () => {
+        const pos: CaretPosition = {
+            lineNumber: 13,
+            column: 25,
+        };
+        const sql = commentOtherLine(syntaxSql, pos.lineNumber);
+        const entities = flink.getAllEntities(sql, pos);
+        expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
+    });
+
+    test('isContainCaret should be falsy if caret position is whitespace after semicolon', () => {
+        const pos: CaretPosition = {
+            lineNumber: 15,
+            column: 32,
+        };
+        const sql = commentOtherLine(syntaxSql, pos.lineNumber);
+        const entities = flink.getAllEntities(sql, pos);
+        expect(entities[0].belongStmt.isContainCaret).toBeFalsy();
+    });
 });

test/parser/hive/suggestion/fixtures/suggestionWithEntity.sql

@@ -20,4 +20,8 @@ INSERT INTO insert_tb PARTITION (country, state) SELECT col1, col2, country, sta
 
 CREATE TABLE IF NOT EXISTS derived_table AS SELECT  FROM origin_table
 
-CREATE TABLE IF NOT EXISTS derived_table AS SELECT id,  FROM origin_table
+CREATE TABLE IF NOT EXISTS derived_table AS SELECT id,  FROM origin_table
+
+SELECT id FROM tb WHERE 
+
+SELECT id FROM tb GROUP BY  ;  

test/parser/hive/suggestion/suggestionWithEntity.test.ts

@@ -118,12 +118,12 @@ describe('Hive SQL Syntax Suggestion with collect entity', () => {
         expect(entities.length).toBe(2);
         expect(entities[0].text).toBe('a');
         expect(entities[0].entityContextType).toBe(EntityContextType.TABLE);
-        expect(entities[0].belongStmt.isContainCaret).toBeFalsy();
+        expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
         expect(entities[0].belongStmt.rootStmt.isContainCaret).toBeTruthy();
 
         expect(entities[1].text).toBe('b');
         expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
-        expect(entities[1].belongStmt.isContainCaret).toBeFalsy();
+        expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
         expect(entities[1].belongStmt.rootStmt.isContainCaret).toBeTruthy();
     });
 
@@ -145,12 +145,12 @@ describe('Hive SQL Syntax Suggestion with collect entity', () => {
         expect(entities.length).toBe(2);
         expect(entities[0].text).toBe('a');
         expect(entities[0].entityContextType).toBe(EntityContextType.TABLE);
-        expect(entities[0].belongStmt.isContainCaret).toBeFalsy();
+        expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
         expect(entities[0].belongStmt.rootStmt.isContainCaret).toBeTruthy();
 
         expect(entities[1].text).toBe('b');
         expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
-        expect(entities[1].belongStmt.isContainCaret).toBeFalsy();
+        expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
         expect(entities[1].belongStmt.rootStmt.isContainCaret).toBeTruthy();
     });
 
@@ -172,12 +172,12 @@ describe('Hive SQL Syntax Suggestion with collect entity', () => {
         expect(entities.length).toBe(2);
         expect(entities[0].text).toBe('page_view_stg');
         expect(entities[0].entityContextType).toBe(EntityContextType.TABLE);
-        expect(entities[0].belongStmt.isContainCaret).toBeFalsy();
+        expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
         expect(entities[0].belongStmt.rootStmt.isContainCaret).toBeTruthy();
 
         expect(entities[1].text).toBe('page_view');
         expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
-        expect(entities[1].belongStmt.isContainCaret).toBeFalsy();
+        expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
         expect(entities[1].belongStmt.rootStmt.isContainCaret).toBeTruthy();
     });
 
@@ -199,12 +199,12 @@ describe('Hive SQL Syntax Suggestion with collect entity', () => {
         expect(entities.length).toBe(2);
         expect(entities[0].text).toBe('page_view_stg');
         expect(entities[0].entityContextType).toBe(EntityContextType.TABLE);
-        expect(entities[0].belongStmt.isContainCaret).toBeFalsy();
+        expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
         expect(entities[0].belongStmt.rootStmt.isContainCaret).toBeTruthy();
 
         expect(entities[1].text).toBe('page_view');
         expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
-        expect(entities[1].belongStmt.isContainCaret).toBeFalsy();
+        expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
         expect(entities[1].belongStmt.rootStmt.isContainCaret).toBeTruthy();
     });
 
@@ -307,4 +307,24 @@ describe('Hive SQL Syntax Suggestion with collect entity', () => {
         expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
         expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
     });
+
+    test('isContainCaret should be truthy if caret position is whitespace at the end of statement', () => {
+        const pos: CaretPosition = {
+            lineNumber: 25,
+            column: 25,
+        };
+        const sql = commentOtherLine(syntaxSql, pos.lineNumber);
+        const entities = hive.getAllEntities(sql, pos);
+        expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
+    });
+
+    test('isContainCaret should be falsy if caret position is whitespace after semicolon', () => {
+        const pos: CaretPosition = {
+            lineNumber: 27,
+            column: 32,
+        };
+        const sql = commentOtherLine(syntaxSql, pos.lineNumber);
+        const entities = hive.getAllEntities(sql, pos);
+        expect(entities[0].belongStmt.isContainCaret).toBeFalsy();
+    });
 });

test/parser/impala/suggestion/fixtures/suggestionWithEntity.sql

@@ -9,3 +9,7 @@ INSERT INTO insert_tb SELECT id,  FROM from_tb;
 CREATE TABLE sorted_census_data AS SELECT  FROM unsorted_census_data;
 
 CREATE TABLE sorted_census_data AS SELECT id,  FROM unsorted_census_data;
+
+SELECT id FROM tb WHERE 
+
+SELECT id FROM tb GROUP BY  ;  

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

@@ -155,4 +155,24 @@ describe('Impala SQL Syntax Suggestion with collect entity', () => {
         expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
         expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
     });
+
+    test('isContainCaret should be truthy if caret position is whitespace at the end of statement', () => {
+        const pos: CaretPosition = {
+            lineNumber: 13,
+            column: 25,
+        };
+        const sql = commentOtherLine(syntaxSql, pos.lineNumber);
+        const entities = impala.getAllEntities(sql, pos);
+        expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
+    });
+
+    test('isContainCaret should be falsy if caret position is whitespace after semicolon', () => {
+        const pos: CaretPosition = {
+            lineNumber: 15,
+            column: 32,
+        };
+        const sql = commentOtherLine(syntaxSql, pos.lineNumber);
+        const entities = impala.getAllEntities(sql, pos);
+        expect(entities[0].belongStmt.isContainCaret).toBeFalsy();
+    });
 });

test/parser/mysql/suggestion/fixtures/suggestionWithEntity.sql

@@ -8,4 +8,8 @@ INSERT INTO insert_tb SELECT id, age,  FROM from_tb;
 
 CREATE TABLE sorted_census_data AS SELECT  FROM unsorted_census_data;
 
-CREATE TABLE sorted_census_data AS SELECT id, age,  FROM unsorted_census_data;
+CREATE TABLE sorted_census_data AS SELECT id, age,  FROM unsorted_census_data;
+
+SELECT id FROM tb WHERE 
+
+SELECT id FROM tb GROUP BY  ;  

test/parser/mysql/suggestion/suggestionWithEntity.test.ts

@@ -153,4 +153,24 @@ describe('MySQL Syntax Suggestion with collect entity', () => {
         expect(entities[1].entityContextType).toBe(EntityContextType.TABLE);
         expect(entities[1].belongStmt.isContainCaret).toBeTruthy();
     });
+
+    test('isContainCaret should be truthy if caret position is whitespace at the end of statement', () => {
+        const pos: CaretPosition = {
+            lineNumber: 13,
+            column: 25,
+        };
+        const sql = commentOtherLine(syntaxSql, pos.lineNumber);
+        const entities = mysql.getAllEntities(sql, pos);
+        expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
+    });
+
+    test('isContainCaret should be falsy if caret position is whitespace after semicolon', () => {
+        const pos: CaretPosition = {
+            lineNumber: 15,
+            column: 32,
+        };
+        const sql = commentOtherLine(syntaxSql, pos.lineNumber);
+        const entities = mysql.getAllEntities(sql, pos);
+        expect(entities[0].belongStmt.isContainCaret).toBeFalsy();
+    });
 });

test/parser/postgresql/suggestion/fixtures/suggestionWithEntity.sql

@@ -10,4 +10,8 @@ CREATE TABLE sorted_census_data AS SELECT  FROM unsorted_census_data;
 
 CREATE TABLE sorted_census_data AS SELECT id, age,  FROM unsorted_census_data;
 
-ALTER TABLE my_table DROP a_column;
+ALTER TABLE my_table DROP a_column;
+
+SELECT id FROM tb WHERE 
+
+SELECT id FROM tb GROUP BY  ;  

test/parser/postgresql/suggestion/suggestionWithEntity.test.ts

@@ -174,4 +174,24 @@ describe('PostgreSql Syntax Suggestion with collect entity', () => {
         expect(entities[0].entityContextType).toBe(EntityContextType.TABLE);
         expect(entities[0].belongStmt?.isContainCaret).toBeTruthy();
     });
+
+    test('isContainCaret should be truthy if caret position is whitespace at the end of statement', () => {
+        const pos: CaretPosition = {
+            lineNumber: 15,
+            column: 25,
+        };
+        const sql = commentOtherLine(syntaxSql, pos.lineNumber);
+        const entities = postgre.getAllEntities(sql, pos);
+        expect(entities[0].belongStmt.isContainCaret).toBeTruthy();
+    });
+
+    test('isContainCaret should be falsy if caret position is whitespace after semicolon', () => {
+        const pos: CaretPosition = {
+            lineNumber: 17,
+            column: 32,
+        };
+        const sql = commentOtherLine(syntaxSql, pos.lineNumber);
+        const entities = postgre.getAllEntities(sql, pos);
+        expect(entities[0].belongStmt.isContainCaret).toBeFalsy();
+    });
 });