diff --git a/README-zh_CN.md b/README-zh_CN.md index 39c12e62..9b7d3b91 100644 --- a/README-zh_CN.md +++ b/README-zh_CN.md @@ -262,6 +262,18 @@ console.log(sqlSlices) [ 'CATALOG', 'FUNCTION', 'TEMPORARY', 'VIEW', 'DATABASE', 'TABLE' ] */ ``` + + 可以通过可选的 `keywordFilter` 移除不需要的关键字候选项。返回 `true` 保留关键字,返回 `false` 移除关键字。 + + ```typescript + const sql = 'SELECT * FROM tb '; + const pos = { lineNumber: 1, column: sql.length + 1 }; + const excludedKeywords = new Set(['WHERE', 'ORDER BY']); + const keywords = flink.getSuggestionAtCaretPosition(sql, pos, { + keywordFilter: (keyword) => !excludedKeywords.has(keyword), + })?.keywords; + ``` + + **获取语法相关自动补全信息** ```typescript import { FlinkSQL } from 'dt-sql-parser'; diff --git a/README.md b/README.md index 37aba125..dd3b5339 100644 --- a/README.md +++ b/README.md @@ -263,6 +263,18 @@ Call the `getAllEntities` method on the SQL instance, pass the SQL content and t [ 'CATALOG', 'FUNCTION', 'TEMPORARY', 'VIEW', 'DATABASE', 'TABLE' ] */ ``` + + Use the optional `keywordFilter` to remove unwanted keyword candidates. Return `true` to keep a keyword and `false` to remove it. + + ```typescript + const sql = 'SELECT * FROM tb '; + const pos = { lineNumber: 1, column: sql.length + 1 }; + const excludedKeywords = new Set(['WHERE', 'ORDER BY']); + const keywords = flink.getSuggestionAtCaretPosition(sql, pos, { + keywordFilter: (keyword) => !excludedKeywords.has(keyword), + })?.keywords; + ``` + + **Obtaining information related to grammar completion** ```javascript import { FlinkSQL } from 'dt-sql-parser'; diff --git a/src/index.ts b/src/index.ts index 5f55a9cf..b7401803 100644 --- a/src/index.ts +++ b/src/index.ts @@ -32,7 +32,12 @@ export { EntityContextType } from './parser/common/types'; export { StmtContextType } from './parser/common/entityCollector'; -export type { CaretPosition, Suggestions, SyntaxSuggestion } from './parser/common/types'; +export type { + CaretPosition, + SuggestionOptions, + Suggestions, + SyntaxSuggestion, +} from './parser/common/types'; export type { WordRange, TextSlice } from './parser/common/textAndWord'; diff --git a/src/parser/common/basicSQL.ts b/src/parser/common/basicSQL.ts index 8d863890..ff59e236 100644 --- a/src/parser/common/basicSQL.ts +++ b/src/parser/common/basicSQL.ts @@ -26,6 +26,7 @@ import { CaretPosition, LOCALE_TYPE, SemanticCollectOptions, + SuggestionOptions, Suggestions, SyntaxSuggestion, } from './types'; @@ -48,6 +49,7 @@ export abstract class BasicSQL< protected _parseTree: PRC | null; protected _parsedInput: string; protected _parseErrors: ParseError[] = []; + private _statementStartTokenTypes: Set | null = null; /** members for cache end */ private _errorListener: ErrorListener = (error) => { @@ -555,15 +557,95 @@ export abstract class BasicSQL< }; } + /** + * Get the minimum statement tree for collecting completion candidates + */ + private getSuggestionParseTree( + parseTree: ParserRuleContext, + caretTokenIndex: number + ): ParserRuleContext { + const children = parseTree.children; + if (!children?.length) return parseTree; + + for (let index = children.length - 1; index >= 0; index--) { + const child = children[index]; + if (!(child instanceof ParserRuleContext)) continue; + + const startTokenIndex = child.start?.tokenIndex; + const stopTokenIndex = child.stop?.tokenIndex; + if ( + startTokenIndex === undefined || + stopTokenIndex === undefined || + startTokenIndex > caretTokenIndex + ) + continue; + + // Use the current statement tree when the caret is inside it + if (stopTokenIndex >= caretTokenIndex) return child; + + // Keep using the current statement until it ends with a semicolon + return child.stop?.text === SQL_SPLIT_SYMBOL_TEXT ? parseTree : child; + } + + return parseTree; + } + + /** + * Collect candidates for the current statement and remove new-statement-only keywords + */ + private collectSuggestionCandidates( + parser: Parser, + parseTree: ParserRuleContext, + caretTokenIndex: number + ): CandidatesCollection { + const core = new CodeCompletionCore(parser); + core.preferredRules = this.preferredRules; + const candidates = core.collectCandidates(caretTokenIndex, parseTree); + const suggestionParseTree = this.getSuggestionParseTree(parseTree, caretTokenIndex); + + if (suggestionParseTree === parseTree) return candidates; + + const statementCore = new CodeCompletionCore(parser); + statementCore.preferredRules = this.preferredRules; + const statementCandidates = statementCore.collectCandidates( + caretTokenIndex, + suggestionParseTree + ); + + if (this._statementStartTokenTypes === null) { + const statementStartCore = new CodeCompletionCore(parser); + statementStartCore.preferredRules = this.preferredRules; + const statementStartCandidates = statementStartCore.collectCandidates(0, parseTree); + this._statementStartTokenTypes = new Set(statementStartCandidates.tokens.keys()); + } + + const tokens = new Map(candidates.tokens); + for (const tokenType of this._statementStartTokenTypes) { + if (!statementCandidates.tokens.has(tokenType)) { + tokens.delete(tokenType); + } else if (tokens.has(tokenType)) { + // Use the current statement follow-list to preserve valid combined keywords + tokens.set(tokenType, statementCandidates.tokens.get(tokenType)!); + } + } + + return { + rules: candidates.rules, + tokens, + }; + } + /** * Get suggestions of syntax and token at caretPosition * @param input source string * @param caretPosition caret position, such as cursor position + * @param options suggestion options * @returns suggestion */ public getSuggestionAtCaretPosition( input: string, - caretPosition: CaretPosition + caretPosition: CaretPosition, + options?: SuggestionOptions ): Suggestions | null { this.parseWithCache(input); if (!this._parseTree) return null; @@ -614,12 +696,11 @@ export abstract class BasicSQL< parseTree = sqlParserIns.program(); } - const core = new CodeCompletionCore(sqlParserIns); - core.preferredRules = this.preferredRules; - // core.showRuleStack = true; - // core.showResult = true; - - const candidates = core.collectCandidates(caretTokenIndex, parseTree); + const candidates = this.collectSuggestionCandidates( + sqlParserIns, + parseTree, + caretTokenIndex + ); const originalSuggestions = this.processCandidates(candidates, allTokens, caretTokenIndex); const syntaxSuggestions: SyntaxSuggestion[] = originalSuggestions.syntax.map( @@ -633,9 +714,14 @@ export abstract class BasicSQL< }; } ); + const keywordFilter = options?.keywordFilter; + const keywords = keywordFilter + ? originalSuggestions.keywords.filter((keyword) => keywordFilter(keyword)) + : originalSuggestions.keywords; + return { syntax: syntaxSuggestions, - keywords: originalSuggestions.keywords, + keywords, }; } diff --git a/src/parser/common/types.ts b/src/parser/common/types.ts index 63bf1cf8..f40cca7d 100644 --- a/src/parser/common/types.ts +++ b/src/parser/common/types.ts @@ -74,6 +74,16 @@ export interface Suggestions { readonly keywords: string[]; } +/** + * Suggested information options + */ +export interface SuggestionOptions { + /** + * Return true to keep the keyword, otherwise remove it + */ + keywordFilter?: (keyword: string) => boolean; +} + export type LOCALE_TYPE = 'zh_CN' | 'en_US'; export interface SemanticContext { diff --git a/test/common/suggestion.test.ts b/test/common/suggestion.test.ts new file mode 100644 index 00000000..4908ca1e --- /dev/null +++ b/test/common/suggestion.test.ts @@ -0,0 +1,184 @@ +import { + FlinkSQL, + GenericSQL, + HiveSQL, + ImpalaSQL, + MySQL, + PostgreSQL, + SparkSQL, + TrinoSQL, +} from 'src/parser'; + +const dialects = [ + ['MySQL', () => new MySQL()], + ['FlinkSQL', () => new FlinkSQL()], + ['SparkSQL', () => new SparkSQL()], + ['HiveSQL', () => new HiveSQL()], + ['PostgreSQL', () => new PostgreSQL()], + ['TrinoSQL', () => new TrinoSQL()], + ['ImpalaSQL', () => new ImpalaSQL()], + ['GenericSQL', () => new GenericSQL()], +] as const; + +describe.each(dialects)('%s suggestion at caret position', (_, createParser) => { + const parser = createParser(); + + test('only suggests keywords for the current statement', () => { + const sql = 'SELECT * FROM t '; + const keywords = parser.getSuggestionAtCaretPosition(sql, { + lineNumber: 1, + column: sql.length + 1, + })?.keywords; + + expect(keywords).toContain('WHERE'); + expect(keywords).not.toContain('SELECT'); + }); + + test('does not suggest a new statement for an incomplete word', () => { + const sql = 'SELECT * FROM t s'; + const keywords = parser.getSuggestionAtCaretPosition(sql, { + lineNumber: 1, + column: sql.length + 1, + })?.keywords; + + expect(keywords).not.toContain('SELECT'); + }); + + test('suggests a new statement after semicolon', () => { + const sql = 'SELECT * FROM t; '; + const keywords = parser.getSuggestionAtCaretPosition(sql, { + lineNumber: 1, + column: sql.length + 1, + })?.keywords; + + expect(keywords).toContain('SELECT'); + }); + + test('only suggests keywords for the second statement', () => { + const sql = 'SELECT * FROM t; SELECT * FROM u '; + const keywords = parser.getSuggestionAtCaretPosition(sql, { + lineNumber: 1, + column: sql.length + 1, + })?.keywords; + + expect(keywords).toContain('WHERE'); + expect(keywords).not.toContain('SELECT'); + }); + + test('filters final keyword suggestions', () => { + const sql = 'SELECT * FROM t; '; + const keywordFilter = jest.fn((keyword: string) => keyword !== 'SELECT'); + const keywords = parser.getSuggestionAtCaretPosition( + sql, + { + lineNumber: 1, + column: sql.length + 1, + }, + { keywordFilter } + )?.keywords; + + expect(keywordFilter).toHaveBeenCalledWith('SELECT'); + expect(keywords).not.toContain('SELECT'); + }); + + test('supports keyword whitelist', () => { + const sql = 'SELECT * FROM t '; + const keywords = parser.getSuggestionAtCaretPosition( + sql, + { + lineNumber: 1, + column: sql.length + 1, + }, + { keywordFilter: (keyword) => keyword === 'WHERE' } + )?.keywords; + + expect(keywords).toEqual(['WHERE']); + }); + + test('filters combined keyword by final text', () => { + const sql = 'SELECT * FROM t '; + const keywords = parser.getSuggestionAtCaretPosition( + sql, + { + lineNumber: 1, + column: sql.length + 1, + }, + { keywordFilter: (keyword) => keyword !== 'ORDER BY' } + )?.keywords; + + expect(keywords).toContain('ORDER'); + expect(keywords).not.toContain('ORDER BY'); + }); + + test('does not filter syntax suggestions', () => { + const sql = 'SELECT * FROM '; + const caretPosition = { + lineNumber: 1, + column: sql.length + 1, + }; + const originalSuggestions = parser.getSuggestionAtCaretPosition(sql, caretPosition); + const filteredSuggestions = parser.getSuggestionAtCaretPosition(sql, caretPosition, { + keywordFilter: () => false, + }); + + expect(originalSuggestions?.syntax.length).toBeGreaterThan(0); + expect(filteredSuggestions?.keywords).toEqual([]); + expect(filteredSuggestions?.syntax).toEqual(originalSuggestions?.syntax); + }); + + test('propagates keyword filter errors', () => { + const sql = 'SELECT * FROM t '; + + expect(() => + parser.getSuggestionAtCaretPosition( + sql, + { + lineNumber: 1, + column: sql.length + 1, + }, + { + keywordFilter: () => { + throw new Error('keyword filter failed'); + }, + } + ) + ).toThrow('keyword filter failed'); + }); +}); + +test('does not add root-absent statement-start keywords to a MySQL nested query', () => { + const parser = new MySQL(); + const sql = 'SELECT * FROM t WHERE EXISTS ('; + const keywords = parser.getSuggestionAtCaretPosition(sql, { + lineNumber: 1, + column: sql.length + 1, + })?.keywords; + + expect(keywords).not.toContain('CREATE'); + expect(keywords).not.toContain('ALTER'); + expect(keywords).not.toContain('DROP'); + expect(keywords).not.toContain('INSERT'); + expect(keywords).not.toContain('CALL'); +}); + +test('keeps MySQL combined keywords from the current statement', () => { + const parser = new MySQL(); + const sql = 'SELECT * FROM t '; + const keywords = parser.getSuggestionAtCaretPosition(sql, { + lineNumber: 1, + column: sql.length + 1, + })?.keywords; + + expect(keywords).toContain('LOCK IN SHARE MODE'); +}); + +test('keeps Impala combined keywords from the current statement', () => { + const parser = new ImpalaSQL(); + const sql = 'CREATE TABLE t (id INT) '; + const keywords = parser.getSuggestionAtCaretPosition(sql, { + lineNumber: 1, + column: sql.length + 1, + })?.keywords; + + expect(keywords).toContain('WITH SERDEPROPERTIES'); +});