Skip to content

Commit

Permalink
feat: add sqlSplitStrategy options
Browse files Browse the repository at this point in the history
  • Loading branch information
JackWang032 committed Jan 2, 2025
1 parent 2f5e9e3 commit 58a65fe
Show file tree
Hide file tree
Showing 24 changed files with 367 additions and 38 deletions.
21 changes: 17 additions & 4 deletions src/parser/common/basicSQL.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -101,7 +107,8 @@ export abstract class BasicSQL<
protected abstract createSemanticContextCollector(
input: string,
caretPosition: CaretPosition,
allTokens: Token[]
allTokens: Token[],
options?: SemanticCollectOptions
): SemanticContextCollector;

/**
Expand Down Expand Up @@ -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);

Expand Down
38 changes: 34 additions & 4 deletions src/parser/common/semanticContextCollector.ts
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -50,6 +64,10 @@ abstract class SemanticContextCollector {
}
}

public readonly options: SemanticCollectOptions = {
sqlSplitStrategy: SqlSplitStrategy.LOOSE,
};

private _tokenIndex: number;
private _allTokens: Token[] = [];

Expand Down Expand Up @@ -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() ||
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
24 changes: 24 additions & 0 deletions src/parser/common/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
13 changes: 10 additions & 3 deletions src/parser/flink/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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(
Expand Down
13 changes: 10 additions & 3 deletions src/parser/hive/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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(
Expand Down
13 changes: 10 additions & 3 deletions src/parser/impala/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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(
Expand Down
13 changes: 10 additions & 3 deletions src/parser/mysql/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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(
Expand Down
13 changes: 10 additions & 3 deletions src/parser/postgresql/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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(
Expand Down
13 changes: 10 additions & 3 deletions src/parser/spark/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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(
Expand Down
13 changes: 10 additions & 3 deletions src/parser/trino/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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([
Expand Down
3 changes: 1 addition & 2 deletions test/parser/flink/contextCollect/fixtures/semantic.sql
Original file line number Diff line number Diff line change
Expand Up @@ -28,5 +28,4 @@ SEL
INSERT INTO t1 VALUES(1);

CREATE TABLE a1(id INT) WITH ('connector' = 'kafka')
CREATE VIEW
INSERT INTO t1 VALUES(1);
CREATE
Loading

0 comments on commit 58a65fe

Please sign in to comment.