fix: #436 improve SQL keyword suggestions

Author: liuxy0551

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';

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';

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';
 

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<number> | 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<WordRange>[] = 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,
         };
     }
 

src/parser/common/types.ts

@@ -74,6 +74,16 @@ export interface Suggestions<T = WordRange> {
     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 {