Feat/follow keywords

package.json

@@ -27,6 +27,7 @@
     "prepublishOnly": "npm run build",
     "antlr4": "node ./scripts/antlr4.js",
     "build": "rm -rf dist && tsc",
+    "watch": "tsc -w",
     "check-types": "tsc -p ./tsconfig.json && tsc -p ./test/tsconfig.json",
     "test": "NODE_OPTIONS=--max_old_space_size=4096 && jest",
     "release": "node ./scripts/release.js",

src/parser/common/tokenUtils.ts

@@ -0,0 +1,52 @@
+/**
+ * Utility function for processing SQL tokens and generating keyword suggestions
+ */
+
+import { Parser } from 'antlr4ng';
+import { CandidatesCollection } from 'antlr4-c3';
+
+/**
+ * Process token candidates and generate a list of keyword suggestions
+ * @param parser SQL parser instance
+ * @param tokens token candidates
+ * @returns list of keyword suggestions
+ */
+export function processTokenCandidates(
+    parser: Parser,
+    tokens: CandidatesCollection['tokens']
+): string[] {
+    const keywords: string[] = [];
+
+    const cleanDisplayName = (displayName: string | null): string => {
+        return displayName && displayName.startsWith("'") && displayName.endsWith("'")
+            ? displayName.slice(1, -1)
+            : displayName || '';
+    };
+
+    const isKeywordToken = (token: number): boolean => {
+        const symbolicName = parser.vocabulary.getSymbolicName(token);
+        return Boolean(symbolicName?.startsWith('KW_'));
+    };
+
+    for (const [token, followSets] of tokens) {
+        const displayName = parser.vocabulary.getDisplayName(token);
+
+        if (!displayName || !isKeywordToken(token)) continue;
+
+        const keyword = cleanDisplayName(displayName);
+        keywords.push(keyword);
+
+        if (followSets.length && followSets.every((s) => isKeywordToken(s))) {
+            const followKeywords = followSets
+                .map((s) => cleanDisplayName(parser.vocabulary.getDisplayName(s)))
+                .filter(Boolean);
+
+            if (followKeywords.length) {
+                const combinedKeyword = [keyword, ...followKeywords].join(' ');
+                keywords.push(combinedKeyword);
+            }
+        }
+    }
+
+    return keywords;
+}

src/parser/flink/index.ts

@@ -1,6 +1,6 @@
 import { CandidatesCollection } from 'antlr4-c3';
 import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
-
+import { processTokenCandidates } from '../common/tokenUtils';
 import { FlinkSqlLexer } from '../../lib/flink/FlinkSqlLexer';
 import { FlinkSqlParser, ProgramContext } from '../../lib/flink/FlinkSqlParser';
 import { BasicSQL } from '../common/basicSQL';
@@ -123,17 +123,9 @@ export class FlinkSQL extends BasicSQL<FlinkSqlLexer, ProgramContext, FlinkSqlPa
             }
         }
 
-        for (let candidate of candidates.tokens) {
-            const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
-            const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
-            if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
-                const keyword =
-                    displayName.startsWith("'") && displayName.endsWith("'")
-                        ? displayName.slice(1, -1)
-                        : displayName;
-                keywords.push(keyword);
-            }
-        }
+        const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
+        keywords.push(...processedKeywords);
+
         return {
             syntax: originalSyntaxSuggestions,
             keywords,

src/parser/hive/index.ts

@@ -1,6 +1,6 @@
 import { CandidatesCollection } from 'antlr4-c3';
 import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
-
+import { processTokenCandidates } from '../common/tokenUtils';
 import { HiveSqlLexer } from '../../lib/hive/HiveSqlLexer';
 import { HiveSqlParser, ProgramContext } from '../../lib/hive/HiveSqlParser';
 import { BasicSQL } from '../common/basicSQL';
@@ -119,17 +119,9 @@ export class HiveSQL extends BasicSQL<HiveSqlLexer, ProgramContext, HiveSqlParse
             }
         }
 
-        for (let candidate of candidates.tokens) {
-            const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
-            const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
-            if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
-                const keyword =
-                    displayName.startsWith("'") && displayName.endsWith("'")
-                        ? displayName.slice(1, -1)
-                        : displayName;
-                keywords.push(keyword);
-            }
-        }
+        const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
+        keywords.push(...processedKeywords);
+
         return {
             syntax: originalSyntaxSuggestions,
             keywords,

src/parser/impala/index.ts

@@ -1,6 +1,6 @@
 import { CandidatesCollection } from 'antlr4-c3';
 import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
-
+import { processTokenCandidates } from '../common/tokenUtils';
 import { ImpalaSqlLexer } from '../../lib/impala/ImpalaSqlLexer';
 import { ImpalaSqlParser, ProgramContext } from '../../lib/impala/ImpalaSqlParser';
 import { BasicSQL } from '../common/basicSQL';
@@ -116,17 +116,9 @@ export class ImpalaSQL extends BasicSQL<ImpalaSqlLexer, ProgramContext, ImpalaSq
             }
         }
 
-        for (let candidate of candidates.tokens) {
-            const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
-            const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
-            if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
-                const keyword =
-                    displayName.startsWith("'") && displayName.endsWith("'")
-                        ? displayName.slice(1, -1)
-                        : displayName;
-                keywords.push(keyword);
-            }
-        }
+        const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
+        keywords.push(...processedKeywords);
+
         return {
             syntax: originalSyntaxSuggestions,
             keywords,

src/parser/mysql/index.ts

@@ -1,6 +1,6 @@
 import { CandidatesCollection } from 'antlr4-c3';
 import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
-
+import { processTokenCandidates } from '../common/tokenUtils';
 import { MySqlLexer } from '../../lib/mysql/MySqlLexer';
 import { MySqlParser, ProgramContext } from '../../lib/mysql/MySqlParser';
 import { BasicSQL } from '../common/basicSQL';
@@ -118,17 +118,8 @@ export class MySQL extends BasicSQL<MySqlLexer, ProgramContext, MySqlParser> {
             }
         }
 
-        for (const candidate of candidates.tokens) {
-            const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
-            const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
-            if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
-                const keyword =
-                    displayName.startsWith("'") && displayName.endsWith("'")
-                        ? displayName.slice(1, -1)
-                        : displayName;
-                keywords.push(keyword);
-            }
-        }
+        const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
+        keywords.push(...processedKeywords);
 
         return {
             syntax: originalSyntaxSuggestions,

src/parser/postgresql/index.ts

@@ -1,5 +1,6 @@
 import { CandidatesCollection } from 'antlr4-c3';
 import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
+import { processTokenCandidates } from '../common/tokenUtils';
 
 import { PostgreSqlLexer } from '../../lib/postgresql/PostgreSqlLexer';
 import { PostgreSqlParser, ProgramContext } from '../../lib/postgresql/PostgreSqlParser';
@@ -137,17 +138,9 @@ export class PostgreSQL extends BasicSQL<PostgreSqlLexer, ProgramContext, Postgr
             }
         }
 
-        for (let candidate of candidates.tokens) {
-            const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
-            const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
-            if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
-                const keyword =
-                    displayName.startsWith("'") && displayName.endsWith("'")
-                        ? displayName.slice(1, -1)
-                        : displayName;
-                keywords.push(keyword);
-            }
-        }
+        const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
+        keywords.push(...processedKeywords);
+
         return {
             syntax: originalSyntaxSuggestions,
             keywords,

src/parser/spark/index.ts

@@ -1,6 +1,6 @@
 import { CandidatesCollection } from 'antlr4-c3';
 import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
-
+import { processTokenCandidates } from '../common/tokenUtils';
 import { SparkSqlLexer } from '../../lib/spark/SparkSqlLexer';
 import { ProgramContext, SparkSqlParser } from '../../lib/spark/SparkSqlParser';
 import { BasicSQL } from '../common/basicSQL';
@@ -118,17 +118,8 @@ export class SparkSQL extends BasicSQL<SparkSqlLexer, ProgramContext, SparkSqlPa
             }
         }
 
-        for (const candidate of candidates.tokens) {
-            const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
-            const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
-            if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
-                const keyword =
-                    displayName.startsWith("'") && displayName.endsWith("'")
-                        ? displayName.slice(1, -1)
-                        : displayName;
-                keywords.push(keyword);
-            }
-        }
+        const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
+        keywords.push(...processedKeywords);
 
         return {
             syntax: originalSyntaxSuggestions,

src/parser/trino/index.ts

@@ -1,6 +1,6 @@
 import { CandidatesCollection } from 'antlr4-c3';
 import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
-
+import { processTokenCandidates } from '../common/tokenUtils';
 import { TrinoSqlLexer } from '../../lib/trino/TrinoSqlLexer';
 import { ProgramContext, TrinoSqlParser } from '../../lib/trino/TrinoSqlParser';
 import { BasicSQL } from '../common/basicSQL';
@@ -128,17 +128,9 @@ export class TrinoSQL extends BasicSQL<TrinoSqlLexer, ProgramContext, TrinoSqlPa
             }
         }
 
-        for (let candidate of candidates.tokens) {
-            const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
-            const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
-            if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
-                const keyword =
-                    displayName.startsWith("'") && displayName.endsWith("'")
-                        ? displayName.slice(1, -1)
-                        : displayName;
-                keywords.push(keyword);
-            }
-        }
+        const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
+        keywords.push(...processedKeywords);
+
         return {
             syntax: originalSyntaxSuggestions,
             keywords,

test/parser/flink/suggestion/fixtures/tokenSuggestion.sql

@@ -4,4 +4,7 @@ USE
 ;
 CREATE 
 ;
-SHOW 
+SHOW 
+;
+CREATE TABLE IF NOT EXISTS
+;

test/parser/flink/suggestion/tokenSuggestion.test.ts

@@ -67,4 +67,30 @@ describe('Flink SQL Token Suggestion', () => {
             'JARS',
         ]);
     });
+
+    test('After CREATE TABLE, show combined keywords', () => {
+        const pos: CaretPosition = {
+            lineNumber: 9,
+            column: 14,
+        };
+        const suggestion = flink.getSuggestionAtCaretPosition(
+            commentOtherLine(tokenSql, pos.lineNumber),
+            pos
+        )?.keywords;
+        expect(suggestion).toContain('IF');
+        expect(suggestion).toContain('IF NOT EXISTS');
+    });
+
+    test('After CREATE TABLE IF, show combined keywords', () => {
+        const pos: CaretPosition = {
+            lineNumber: 9,
+            column: 17,
+        };
+        const suggestion = flink.getSuggestionAtCaretPosition(
+            commentOtherLine(tokenSql, pos.lineNumber),
+            pos
+        )?.keywords;
+        expect(suggestion).toContain('NOT');
+        expect(suggestion).toContain('NOT EXISTS');
+    });
 });

test/parser/hive/suggestion/fixtures/tokenSuggestion.sql

@@ -18,3 +18,5 @@ LOAD
 ;
 SHOW 
 ;
+CREATE TABLE IF NOT EXISTS
+;

test/parser/hive/suggestion/tokenSuggestion.test.ts

@@ -33,6 +33,9 @@ describe('Hive SQL Token Suggestion', () => {
             'MATERIALIZED',
             'VIEW',
             'TABLE',
+            'RESOURCE PLAN',
+            'SCHEDULED QUERY',
+            'MATERIALIZED VIEW',
         ]);
     });
 
@@ -68,6 +71,11 @@ describe('Hive SQL Token Suggestion', () => {
             'REMOTE',
             'DATABASE',
             'SCHEMA',
+            'RESOURCE PLAN',
+            'SCHEDULED QUERY',
+            'MATERIALIZED VIEW',
+            'OR REPLACE',
+            'MANAGED TABLE',
         ]);
     });
 
@@ -129,6 +137,9 @@ describe('Hive SQL Token Suggestion', () => {
             'TABLE',
             'DATABASE',
             'SCHEMA',
+            'RESOURCE PLAN',
+            'MATERIALIZED VIEW',
+            'SCHEDULED QUERY',
         ]);
     });
 
@@ -217,6 +228,36 @@ describe('Hive SQL Token Suggestion', () => {
             'EXTENDED',
             'DATABASES',
             'SCHEMAS',
+            'CURRENT ROLES',
+            'ROLE GRANT',
+            'TABLE EXTENDED',
+            'MATERIALIZED VIEWS',
         ]);
     });
+
+    test('After CREATE TABLE, show combined keywords', () => {
+        const pos: CaretPosition = {
+            lineNumber: 21,
+            column: 14,
+        };
+        const suggestion = hive.getSuggestionAtCaretPosition(
+            commentOtherLine(tokenSql, pos.lineNumber),
+            pos
+        )?.keywords;
+        expect(suggestion).toContain('IF');
+        expect(suggestion).toContain('IF NOT EXISTS');
+    });
+
+    test('After CREATE TABLE IF, show combined keywords', () => {
+        const pos: CaretPosition = {
+            lineNumber: 21,
+            column: 17,
+        };
+        const suggestion = hive.getSuggestionAtCaretPosition(
+            commentOtherLine(tokenSql, pos.lineNumber),
+            pos
+        )?.keywords;
+        expect(suggestion).toContain('NOT');
+        expect(suggestion).toContain('NOT EXISTS');
+    });
 });

test/parser/impala/suggestion/fixtures/tokenSuggestion.sql

@@ -9,3 +9,5 @@ INSERT ;
 SHOW ;
 
 CREATE TABLE t1 (id );
+
+CREATE TABLE IF NOT EXISTS;