feat: add sqlSplitStrategy options

Author: JackWang032

src/parser/common/basicSQL.ts

@@ -14,7 +14,13 @@ import { CandidatesCollection, CodeCompletionCore } from 'antlr4-c3';
 import { SQLParserBase } from '../../lib/SQLParserBase';
 import { findCaretTokenIndex } from './findCaretTokenIndex';
 import { ctxToText, tokenToWord, WordRange, TextSlice } from './textAndWord';
-import { CaretPosition, LOCALE_TYPE, Suggestions, SyntaxSuggestion } from './types';
+import {
+    CaretPosition,
+    LOCALE_TYPE,
+    SemanticCollectOptions,
+    Suggestions,
+    SyntaxSuggestion,
+} from './types';
 import { ParseError, ErrorListener } from './parseErrorListener';
 import { ErrorStrategy } from './errorStrategy';
 import type { SplitListener } from './splitListener';
@@ -101,7 +107,8 @@ export abstract class BasicSQL<
     protected abstract createSemanticContextCollector(
         input: string,
         caretPosition: CaretPosition,
-        allTokens: Token[]
+        allTokens: Token[],
+        options?: SemanticCollectOptions
     ): SemanticContextCollector;
 
     /**
@@ -466,15 +473,21 @@ export abstract class BasicSQL<
      * Get semantic context infos
      * @param input source string
      * @param caretPosition caret position, such as cursor position
+     * @param options semantic context options
      * @returns analyzed semantic context
      */
-    public getSemanticContextAtCaretPosition(input: string, caretPosition: CaretPosition) {
+    public getSemanticContextAtCaretPosition(
+        input: string,
+        caretPosition: CaretPosition,
+        options?: SemanticCollectOptions
+    ) {
         const allTokens = this.getAllTokens(input);
         const parseTree = this.parseWithCache(input);
         const statementContextListener = this.createSemanticContextCollector(
             input,
             caretPosition,
-            allTokens
+            allTokens,
+            options
         );
         this.listen(statementContextListener, parseTree);
 

src/parser/common/semanticContextCollector.ts

@@ -1,18 +1,32 @@
 import { ErrorNode, ParserRuleContext, TerminalNode, Token } from 'antlr4ng';
 import { findCaretTokenIndex } from '../common/findCaretTokenIndex';
-import { CaretPosition, SemanticContext } from '../common/types';
+import {
+    CaretPosition,
+    SemanticCollectOptions,
+    SemanticContext,
+    SqlSplitStrategy,
+} from '../common/types';
 
 export const SQL_SPLIT_SYMBOL_TEXT = ';';
 
 abstract class SemanticContextCollector {
-    constructor(_input: string, caretPosition: CaretPosition, allTokens: Token[]) {
+    constructor(
+        _input: string,
+        caretPosition: CaretPosition,
+        allTokens: Token[],
+        options?: SemanticCollectOptions
+    ) {
         // If caretPosition token is whiteSpace, tokenIndex may be undefined.
         const tokenIndex = findCaretTokenIndex(caretPosition, allTokens);
 
         if (tokenIndex !== undefined) {
             this._tokenIndex = tokenIndex;
         }
         this._allTokens = allTokens;
+        this.options = {
+            ...this.options,
+            ...options,
+        };
 
         if (allTokens?.length) {
             let i = tokenIndex ? tokenIndex - 1 : allTokens.length - 1;
@@ -50,6 +64,10 @@ abstract class SemanticContextCollector {
         }
     }
 
+    public readonly options: SemanticCollectOptions = {
+        sqlSplitStrategy: SqlSplitStrategy.LOOSE,
+    };
+
     private _tokenIndex: number;
     private _allTokens: Token[] = [];
 
@@ -117,6 +135,8 @@ abstract class SemanticContextCollector {
      * It should be called in each language's own `enterStatement`.
      */
     protected visitStatement(ctx: ParserRuleContext) {
+        if (this.options.sqlSplitStrategy === SqlSplitStrategy.STRICT) return;
+
         const isWhiteSpaceToken =
             this._tokenIndex === undefined ||
             this._allTokens[this._tokenIndex]?.type === this.getWhiteSpaceRuleType() ||
@@ -135,7 +155,12 @@ abstract class SemanticContextCollector {
      * Uncomplete keyword will be error node
      */
     visitErrorNode(node: ErrorNode): void {
-        if (node.symbol.tokenIndex !== this._tokenIndex || this._isNewStatement) return;
+        if (
+            node.symbol.tokenIndex !== this._tokenIndex ||
+            this._isNewStatement ||
+            this.options.sqlSplitStrategy === SqlSplitStrategy.STRICT
+        )
+            return;
 
         let parent: ParserRuleContext | null = node.parent as ParserRuleContext;
         let currentNode: TerminalNode | ParserRuleContext = node;
@@ -188,7 +213,12 @@ abstract class SemanticContextCollector {
     }
 
     visitTerminal(node: TerminalNode): void {
-        if (node.symbol.tokenIndex !== this._tokenIndex || this._isNewStatement) return;
+        if (
+            node.symbol.tokenIndex !== this._tokenIndex ||
+            this._isNewStatement ||
+            this.options.sqlSplitStrategy === SqlSplitStrategy.STRICT
+        )
+            return;
 
         let currentNode: TerminalNode | ParserRuleContext = node;
         let parent = node.parent as ParserRuleContext | null;

src/parser/common/types.ts

@@ -73,3 +73,27 @@ export type LOCALE_TYPE = 'zh_CN' | 'en_US';
 export interface SemanticContext {
     isNewStatement: boolean;
 }
+
+export enum SqlSplitStrategy {
+    /** Only end the statement with semicolon symbol */
+    STRICT,
+    /** Based on parse tree to split statements */
+    LOOSE,
+}
+
+export interface SemanticCollectOptions {
+    /**
+     * `sqlSplitStrategy` will affects the result of `isNewStatement`;
+     *
+     * For example:
+     *
+     * The sql is "select id from t1   create\<cart_position\>"
+     *
+     * - `SqlSplitStrategy.STRICT`: split symbol `;` is missing after select statement so that it considerd as one statement, and `isNewStatement` is false
+     *
+     * - `SqlSplitStrategy.LOOSE`: in parse tree, it will parse to "select id from t1" and "create" two single statement, so `isNewStatement` is true
+     *
+     * @default SqlSplitStrategy.STRICT
+     */
+    sqlSplitStrategy?: SqlSplitStrategy;
+}

src/parser/flink/index.ts

@@ -2,7 +2,13 @@ import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
 import { CandidatesCollection } from 'antlr4-c3';
 import { FlinkSqlLexer } from '../../lib/flink/FlinkSqlLexer';
 import { FlinkSqlParser, ProgramContext } from '../../lib/flink/FlinkSqlParser';
-import { CaretPosition, EntityContextType, Suggestions, SyntaxSuggestion } from '../common/types';
+import {
+    CaretPosition,
+    EntityContextType,
+    SemanticCollectOptions,
+    Suggestions,
+    SyntaxSuggestion,
+} from '../common/types';
 import { BasicSQL } from '../common/basicSQL';
 import { StmtContextType } from '../common/entityCollector';
 import { FlinkSqlSplitListener } from './flinkSplitListener';
@@ -52,9 +58,10 @@ export class FlinkSQL extends BasicSQL<FlinkSqlLexer, ProgramContext, FlinkSqlPa
     protected createSemanticContextCollector(
         input: string,
         caretPosition: CaretPosition,
-        allTokens: Token[]
+        allTokens: Token[],
+        options?: SemanticCollectOptions
     ) {
-        return new FlinkSemanticContextCollector(input, caretPosition, allTokens);
+        return new FlinkSemanticContextCollector(input, caretPosition, allTokens, options);
     }
 
     protected processCandidates(

src/parser/hive/index.ts

@@ -4,7 +4,13 @@ import { HiveSqlLexer } from '../../lib/hive/HiveSqlLexer';
 import { HiveSqlParser, ProgramContext } from '../../lib/hive/HiveSqlParser';
 import { BasicSQL } from '../common/basicSQL';
 
-import { CaretPosition, EntityContextType, Suggestions, SyntaxSuggestion } from '../common/types';
+import {
+    CaretPosition,
+    EntityContextType,
+    SemanticCollectOptions,
+    Suggestions,
+    SyntaxSuggestion,
+} from '../common/types';
 import { StmtContextType } from '../common/entityCollector';
 import { HiveSqlSplitListener } from './hiveSplitListener';
 import { HiveEntityCollector } from './hiveEntityCollector';
@@ -53,9 +59,10 @@ export class HiveSQL extends BasicSQL<HiveSqlLexer, ProgramContext, HiveSqlParse
     protected createSemanticContextCollector(
         input: string,
         caretPosition: CaretPosition,
-        allTokens: Token[]
+        allTokens: Token[],
+        options?: SemanticCollectOptions
     ) {
-        return new HiveSemanticContextCollector(input, caretPosition, allTokens);
+        return new HiveSemanticContextCollector(input, caretPosition, allTokens, options);
     }
 
     protected processCandidates(

src/parser/impala/index.ts

@@ -3,7 +3,13 @@ import { CandidatesCollection } from 'antlr4-c3';
 import { ImpalaSqlLexer } from '../../lib/impala/ImpalaSqlLexer';
 import { ImpalaSqlParser, ProgramContext } from '../../lib/impala/ImpalaSqlParser';
 import { BasicSQL } from '../common/basicSQL';
-import { CaretPosition, EntityContextType, Suggestions, SyntaxSuggestion } from '../common/types';
+import {
+    CaretPosition,
+    EntityContextType,
+    SemanticCollectOptions,
+    Suggestions,
+    SyntaxSuggestion,
+} from '../common/types';
 import { StmtContextType } from '../common/entityCollector';
 import { ImpalaSqlSplitListener } from './impalaSplitListener';
 import { ImpalaEntityCollector } from './impalaEntityCollector';
@@ -51,9 +57,10 @@ export class ImpalaSQL extends BasicSQL<ImpalaSqlLexer, ProgramContext, ImpalaSq
     protected createSemanticContextCollector(
         input: string,
         caretPosition: CaretPosition,
-        allTokens: Token[]
+        allTokens: Token[],
+        options?: SemanticCollectOptions
     ) {
-        return new ImpalaSemanticContextCollector(input, caretPosition, allTokens);
+        return new ImpalaSemanticContextCollector(input, caretPosition, allTokens, options);
     }
 
     protected processCandidates(

src/parser/mysql/index.ts

@@ -3,7 +3,13 @@ import { CandidatesCollection } from 'antlr4-c3';
 import { MySqlLexer } from '../../lib/mysql/MySqlLexer';
 import { MySqlParser, ProgramContext } from '../../lib/mysql/MySqlParser';
 import { BasicSQL } from '../common/basicSQL';
-import { Suggestions, EntityContextType, SyntaxSuggestion, CaretPosition } from '../common/types';
+import {
+    Suggestions,
+    EntityContextType,
+    SyntaxSuggestion,
+    CaretPosition,
+    SemanticCollectOptions,
+} from '../common/types';
 import { StmtContextType } from '../common/entityCollector';
 import { MysqlSplitListener } from './mysqlSplitListener';
 import { MySqlEntityCollector } from './mysqlEntityCollector';
@@ -51,9 +57,10 @@ export class MySQL extends BasicSQL<MySqlLexer, ProgramContext, MySqlParser> {
     protected createSemanticContextCollector(
         input: string,
         caretPosition: CaretPosition,
-        allTokens: Token[]
+        allTokens: Token[],
+        options?: SemanticCollectOptions
     ) {
-        return new MySqlSemanticContextCollector(input, caretPosition, allTokens);
+        return new MySqlSemanticContextCollector(input, caretPosition, allTokens, options);
     }
 
     protected processCandidates(

src/parser/postgresql/index.ts

@@ -3,7 +3,13 @@ import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
 
 import { PostgreSqlLexer } from '../../lib/postgresql/PostgreSqlLexer';
 import { PostgreSqlParser, ProgramContext } from '../../lib/postgresql/PostgreSqlParser';
-import { CaretPosition, EntityContextType, Suggestions, SyntaxSuggestion } from '../common/types';
+import {
+    CaretPosition,
+    EntityContextType,
+    SemanticCollectOptions,
+    Suggestions,
+    SyntaxSuggestion,
+} from '../common/types';
 import { BasicSQL } from '../common/basicSQL';
 import { StmtContextType } from '../common/entityCollector';
 import { PostgreSqlEntityCollector } from './postgreEntityCollector';
@@ -56,9 +62,10 @@ export class PostgreSQL extends BasicSQL<PostgreSqlLexer, ProgramContext, Postgr
     protected createSemanticContextCollector(
         input: string,
         caretPosition: CaretPosition,
-        allTokens: Token[]
+        allTokens: Token[],
+        options?: SemanticCollectOptions
     ) {
-        return new PostgreSemanticContextCollector(input, caretPosition, allTokens);
+        return new PostgreSemanticContextCollector(input, caretPosition, allTokens, options);
     }
 
     protected processCandidates(

src/parser/spark/index.ts

@@ -3,7 +3,13 @@ import { CandidatesCollection } from 'antlr4-c3';
 import { SparkSqlLexer } from '../../lib/spark/SparkSqlLexer';
 import { SparkSqlParser, ProgramContext } from '../../lib/spark/SparkSqlParser';
 import { BasicSQL } from '../common/basicSQL';
-import { Suggestions, EntityContextType, SyntaxSuggestion, CaretPosition } from '../common/types';
+import {
+    Suggestions,
+    EntityContextType,
+    SyntaxSuggestion,
+    CaretPosition,
+    SemanticCollectOptions,
+} from '../common/types';
 import { StmtContextType } from '../common/entityCollector';
 import { SparkSqlSplitListener } from './sparkSplitListener';
 import { SparkEntityCollector } from './sparkEntityCollector';
@@ -51,9 +57,10 @@ export class SparkSQL extends BasicSQL<SparkSqlLexer, ProgramContext, SparkSqlPa
     protected createSemanticContextCollector(
         input: string,
         caretPosition: CaretPosition,
-        allTokens: Token[]
+        allTokens: Token[],
+        options?: SemanticCollectOptions
     ) {
-        return new SparkSemanticContextCollector(input, caretPosition, allTokens);
+        return new SparkSemanticContextCollector(input, caretPosition, allTokens, options);
     }
 
     protected processCandidates(

src/parser/trino/index.ts

@@ -3,7 +3,13 @@ import { CandidatesCollection } from 'antlr4-c3';
 import { TrinoSqlLexer } from '../../lib/trino/TrinoSqlLexer';
 import { TrinoSqlParser, ProgramContext } from '../../lib/trino/TrinoSqlParser';
 import { BasicSQL } from '../common/basicSQL';
-import { Suggestions, EntityContextType, SyntaxSuggestion, CaretPosition } from '../common/types';
+import {
+    Suggestions,
+    EntityContextType,
+    SyntaxSuggestion,
+    CaretPosition,
+    SemanticCollectOptions,
+} from '../common/types';
 import { StmtContextType } from '../common/entityCollector';
 import { TrinoSqlSplitListener } from './trinoSplitListener';
 import { TrinoEntityCollector } from './trinoEntityCollector';
@@ -38,9 +44,10 @@ export class TrinoSQL extends BasicSQL<TrinoSqlLexer, ProgramContext, TrinoSqlPa
     protected createSemanticContextCollector(
         input: string,
         caretPosition: CaretPosition,
-        allTokens: Token[]
+        allTokens: Token[],
+        options?: SemanticCollectOptions
     ) {
-        return new TrinoSemanticContextCollector(input, caretPosition, allTokens);
+        return new TrinoSemanticContextCollector(input, caretPosition, allTokens, options);
     }
 
     protected preferredRules: Set<number> = new Set([