fix: update isContainCaret judgment when caret position token is whit… (#390)

* fix: update isContainCaret judgment when caret position token is whitespace

* fix: remove unnecessary +Infinity

Author: JackWang032

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  ;