From 9125e8b2d6f95357949fbc0afe3a70a8a8335d0f Mon Sep 17 00:00:00 2001 From: JackWang032 <2522134117@qq.com> Date: Thu, 20 Mar 2025 14:56:46 +0800 Subject: [PATCH 1/2] feat: provide follow keywords when get suggestions --- src/parser/common/tokenUtils.ts | 52 +++++++++++++++++++ src/parser/flink/index.ts | 16 ++---- src/parser/hive/index.ts | 16 ++---- src/parser/impala/index.ts | 16 ++---- src/parser/mysql/index.ts | 15 ++---- src/parser/postgresql/index.ts | 15 ++---- src/parser/spark/index.ts | 15 ++---- src/parser/trino/index.ts | 16 ++---- .../suggestion/fixtures/tokenSuggestion.sql | 5 +- .../flink/suggestion/tokenSuggestion.test.ts | 26 ++++++++++ .../suggestion/fixtures/tokenSuggestion.sql | 2 + .../hive/suggestion/tokenSuggestion.test.ts | 41 +++++++++++++++ .../suggestion/fixtures/tokenSuggestion.sql | 2 + .../impala/suggestion/tokenSuggestion.test.ts | 30 +++++++++++ .../suggestion/fixtures/tokenSuggestion.sql | 2 + .../mysql/suggestion/tokenSuggestion.test.ts | 45 +++++++++++++++- .../suggestion/fixtures/tokenSuggestion.sql | 3 +- .../suggestion/tokenSuggestion.test.ts | 40 ++++++++++++++ .../suggestion/fixtures/tokenSuggestion.sql | 2 + .../spark/suggestion/tokenSuggestion.test.ts | 31 +++++++++++ .../suggestion/fixtures/tokenSuggestion.sql | 2 + .../trino/suggestion/tokenSuggestion.test.ts | 37 ++++++++++++- 22 files changed, 341 insertions(+), 88 deletions(-) create mode 100644 src/parser/common/tokenUtils.ts diff --git a/src/parser/common/tokenUtils.ts b/src/parser/common/tokenUtils.ts new file mode 100644 index 00000000..2f9394da --- /dev/null +++ b/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; +} diff --git a/src/parser/flink/index.ts b/src/parser/flink/index.ts index b373b736..17c6af3d 100644 --- a/src/parser/flink/index.ts +++ b/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 { } } - 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, diff --git a/src/parser/postgresql/index.ts b/src/parser/postgresql/index.ts index d655a280..77961092 100644 --- a/src/parser/postgresql/index.ts +++ b/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 { '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'); + }); }); diff --git a/test/parser/hive/suggestion/fixtures/tokenSuggestion.sql b/test/parser/hive/suggestion/fixtures/tokenSuggestion.sql index 87d3c7f4..ab341189 100644 --- a/test/parser/hive/suggestion/fixtures/tokenSuggestion.sql +++ b/test/parser/hive/suggestion/fixtures/tokenSuggestion.sql @@ -18,3 +18,5 @@ LOAD ; SHOW ; +CREATE TABLE IF NOT EXISTS +; diff --git a/test/parser/hive/suggestion/tokenSuggestion.test.ts b/test/parser/hive/suggestion/tokenSuggestion.test.ts index 06c0e3d9..d74068e2 100644 --- a/test/parser/hive/suggestion/tokenSuggestion.test.ts +++ b/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'); + }); }); diff --git a/test/parser/impala/suggestion/fixtures/tokenSuggestion.sql b/test/parser/impala/suggestion/fixtures/tokenSuggestion.sql index 19855565..156afaea 100644 --- a/test/parser/impala/suggestion/fixtures/tokenSuggestion.sql +++ b/test/parser/impala/suggestion/fixtures/tokenSuggestion.sql @@ -9,3 +9,5 @@ INSERT ; SHOW ; CREATE TABLE t1 (id ); + +CREATE TABLE IF NOT EXISTS; \ No newline at end of file diff --git a/test/parser/impala/suggestion/tokenSuggestion.test.ts b/test/parser/impala/suggestion/tokenSuggestion.test.ts index 627d7293..93249a91 100644 --- a/test/parser/impala/suggestion/tokenSuggestion.test.ts +++ b/test/parser/impala/suggestion/tokenSuggestion.test.ts @@ -107,6 +107,10 @@ describe('Impala SQL Token Suggestion', () => { 'ROLE', 'ROLES', 'CURRENT', + 'TABLE STATS', + 'COLUMN STATS', + 'FILES IN', + 'ROLE GRANT GROUP', ]); }); @@ -143,4 +147,30 @@ describe('Impala SQL Token Suggestion', () => { expect(dataTypes.every((dataType) => suggestion.includes(dataType))).toBe(true); }); + + test('After CREATE TABLE, show combined keywords', () => { + const pos: CaretPosition = { + lineNumber: 13, + column: 14, + }; + const suggestion = impala.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: 13, + column: 17, + }; + const suggestion = impala.getSuggestionAtCaretPosition( + commentOtherLine(tokenSql, pos.lineNumber), + pos + )?.keywords; + expect(suggestion).toContain('NOT'); + expect(suggestion).toContain('NOT EXISTS'); + }); }); diff --git a/test/parser/mysql/suggestion/fixtures/tokenSuggestion.sql b/test/parser/mysql/suggestion/fixtures/tokenSuggestion.sql index 11f02aba..0733f754 100644 --- a/test/parser/mysql/suggestion/fixtures/tokenSuggestion.sql +++ b/test/parser/mysql/suggestion/fixtures/tokenSuggestion.sql @@ -14,3 +14,5 @@ LOAD ; SHOW ; +CREATE TABLE IF NOT EXISTS +; \ No newline at end of file diff --git a/test/parser/mysql/suggestion/tokenSuggestion.test.ts b/test/parser/mysql/suggestion/tokenSuggestion.test.ts index 2c913cff..9d3ad330 100644 --- a/test/parser/mysql/suggestion/tokenSuggestion.test.ts +++ b/test/parser/mysql/suggestion/tokenSuggestion.test.ts @@ -37,6 +37,10 @@ describe('MySQL Token Suggestion', () => { 'EVENT', 'DATABASE', 'SCHEMA', + 'RESOURCE GROUP', + 'SQL SECURITY', + 'LOGFILE GROUP', + 'INSTANCE ROTATE INNODB MASTER KEY', ]); }); @@ -79,6 +83,11 @@ describe('MySQL Token Suggestion', () => { 'EVENT', 'DATABASE', 'SCHEMA', + 'RESOURCE GROUP', + 'SQL SECURITY', + 'OR REPLACE', + 'IF NOT EXISTS', + 'LOGFILE GROUP', ]); }); @@ -116,6 +125,7 @@ describe('MySQL Token Suggestion', () => { 'FORMAT', 'PARTITIONS', 'EXTENDED', + 'FOR CONNECTION', ]); }); @@ -149,6 +159,9 @@ describe('MySQL Token Suggestion', () => { 'EVENT', 'DATABASE', 'SCHEMA', + 'RESOURCE GROUP', + 'SPATIAL REFERENCE SYSTEM', + 'LOGFILE GROUP', ]); }); @@ -181,7 +194,7 @@ describe('MySQL Token Suggestion', () => { pos )?.keywords; - expect(suggestion).toMatchUnorderedArray(['INDEX', 'XML', 'DATA']); + expect(suggestion).toMatchUnorderedArray(['INDEX', 'XML', 'DATA', 'INDEX INTO CACHE']); }); test('After SHOW', () => { @@ -240,6 +253,36 @@ describe('MySQL Token Suggestion', () => { 'BINLOG', 'RELAYLOG', 'BINARY', + 'OPEN TABLES', + 'TABLE STATUS', + 'CHARACTER SET', ]); }); + + test('After CREATE TABLE, show combined keywords', () => { + const pos: CaretPosition = { + lineNumber: 17, + column: 14, + }; + const suggestion = mysql.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: 17, + column: 17, + }; + const suggestion = mysql.getSuggestionAtCaretPosition( + commentOtherLine(tokenSql, pos.lineNumber), + pos + )?.keywords; + expect(suggestion).toContain('NOT'); + expect(suggestion).toContain('NOT EXISTS'); + }); }); diff --git a/test/parser/postgresql/suggestion/fixtures/tokenSuggestion.sql b/test/parser/postgresql/suggestion/fixtures/tokenSuggestion.sql index 6eaefbcf..f49facfc 100644 --- a/test/parser/postgresql/suggestion/fixtures/tokenSuggestion.sql +++ b/test/parser/postgresql/suggestion/fixtures/tokenSuggestion.sql @@ -8,5 +8,4 @@ DELETE ; CREATE ; - - +CREATE TABLE IF NOT EXISTS; \ No newline at end of file diff --git a/test/parser/postgresql/suggestion/tokenSuggestion.test.ts b/test/parser/postgresql/suggestion/tokenSuggestion.test.ts index 32bd8479..18625ca0 100644 --- a/test/parser/postgresql/suggestion/tokenSuggestion.test.ts +++ b/test/parser/postgresql/suggestion/tokenSuggestion.test.ts @@ -54,6 +54,10 @@ describe('Postgres SQL Token Suggestion', () => { 'LARGE', 'EXTENSION', 'DEFAULT', + 'TEXT SEARCH', + 'EVENT TRIGGER', + 'LARGE OBJECT', + 'DEFAULT PRIVILEGES', ]); }); @@ -114,6 +118,12 @@ describe('Postgres SQL Token Suggestion', () => { 'CAST', 'ASSERTION', 'ACCESS', + 'RECURSIVE VIEW', + 'TEXT SEARCH', + 'EVENT TRIGGER', + 'TRANSFORM FOR', + 'MATERIALIZED VIEW', + 'ACCESS METHOD', ]); }); @@ -176,6 +186,10 @@ describe('Postgres SQL Token Suggestion', () => { 'MATERIALIZED', 'SEQUENCE', 'TABLE', + 'OWNED BY', + 'TEXT SEARCH', + 'EVENT TRIGGER', + 'ACCESS METHOD', ]); }); @@ -190,4 +204,30 @@ describe('Postgres SQL Token Suggestion', () => { )?.keywords; expect(suggestion).toMatchUnorderedArray(['INTO']); }); + + test('After CREATE TABLE, show combined keywords', () => { + const pos: CaretPosition = { + lineNumber: 11, + column: 14, + }; + const suggestion = postgresql.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: 11, + column: 17, + }; + const suggestion = postgresql.getSuggestionAtCaretPosition( + commentOtherLine(tokenSql, pos.lineNumber), + pos + )?.keywords; + expect(suggestion).toContain('NOT'); + expect(suggestion).toContain('NOT EXISTS'); + }); }); diff --git a/test/parser/spark/suggestion/fixtures/tokenSuggestion.sql b/test/parser/spark/suggestion/fixtures/tokenSuggestion.sql index 64741647..cdb60e02 100644 --- a/test/parser/spark/suggestion/fixtures/tokenSuggestion.sql +++ b/test/parser/spark/suggestion/fixtures/tokenSuggestion.sql @@ -16,3 +16,5 @@ SHOW ; EXPORT ; +CREATE TABLE IF NOT EXISTS +; diff --git a/test/parser/spark/suggestion/tokenSuggestion.test.ts b/test/parser/spark/suggestion/tokenSuggestion.test.ts index 4cfd9eda..f12e647d 100644 --- a/test/parser/spark/suggestion/tokenSuggestion.test.ts +++ b/test/parser/spark/suggestion/tokenSuggestion.test.ts @@ -27,6 +27,7 @@ describe('Spark SQL Token Suggestion', () => { 'DATABASE', 'NAMESPACE', 'SCHEMA', + 'MATERIALIZED VIEW', ]); }); @@ -54,6 +55,7 @@ describe('Spark SQL Token Suggestion', () => { 'DATABASE', 'NAMESPACE', 'SCHEMA', + 'MATERIALIZED VIEW', ]); }); @@ -117,6 +119,7 @@ describe('Spark SQL Token Suggestion', () => { 'DATABASE', 'NAMESPACE', 'SCHEMA', + 'MATERIALIZED VIEW', ]); }); @@ -182,6 +185,8 @@ describe('Spark SQL Token Suggestion', () => { 'DATABASES', 'NAMESPACES', 'SCHEMAS', + 'MATERIALIZED VIEWS', + 'TABLE EXTENDED', ]); }); @@ -197,4 +202,30 @@ describe('Spark SQL Token Suggestion', () => { expect(suggestion).toMatchUnorderedArray(['TABLE']); }); + + test('After CREATE TABLE, show combined keywords', () => { + const pos: CaretPosition = { + lineNumber: 19, + column: 14, + }; + const suggestion = spark.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: 19, + column: 17, + }; + const suggestion = spark.getSuggestionAtCaretPosition( + commentOtherLine(tokenSql, pos.lineNumber), + pos + )?.keywords; + expect(suggestion).toContain('NOT'); + expect(suggestion).toContain('NOT EXISTS'); + }); }); diff --git a/test/parser/trino/suggestion/fixtures/tokenSuggestion.sql b/test/parser/trino/suggestion/fixtures/tokenSuggestion.sql index 9ff940d7..e1d7328a 100644 --- a/test/parser/trino/suggestion/fixtures/tokenSuggestion.sql +++ b/test/parser/trino/suggestion/fixtures/tokenSuggestion.sql @@ -11,3 +11,5 @@ DESCRIBE ; DROP ; INSERT ; + +CREATE TABLE IF NOT EXISTS ; \ No newline at end of file diff --git a/test/parser/trino/suggestion/tokenSuggestion.test.ts b/test/parser/trino/suggestion/tokenSuggestion.test.ts index d7398dc2..239a4bb1 100644 --- a/test/parser/trino/suggestion/tokenSuggestion.test.ts +++ b/test/parser/trino/suggestion/tokenSuggestion.test.ts @@ -19,7 +19,13 @@ describe('Trino SQL Token Suggestion', () => { pos )?.keywords; - expect(suggestion).toMatchUnorderedArray(['VIEW', 'MATERIALIZED', 'TABLE', 'SCHEMA']); + expect(suggestion).toMatchUnorderedArray([ + 'VIEW', + 'MATERIALIZED', + 'TABLE', + 'SCHEMA', + 'MATERIALIZED VIEW', + ]); }); test('After CREATE', () => { @@ -41,6 +47,8 @@ describe('Trino SQL Token Suggestion', () => { 'TABLE', 'SCHEMA', 'CATALOG', + 'OR REPLACE', + 'MATERIALIZED VIEW', ]); }); @@ -109,6 +117,7 @@ describe('Trino SQL Token Suggestion', () => { 'TABLE', 'SCHEMA', 'CATALOG', + 'MATERIALIZED VIEW', ]); }); @@ -124,4 +133,30 @@ describe('Trino SQL Token Suggestion', () => { expect(suggestion).toMatchUnorderedArray(['INTO']); }); + + test('After CREATE TABLE, show combined keywords', () => { + const pos: CaretPosition = { + lineNumber: 15, + column: 14, + }; + const suggestion = trino.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: 15, + column: 17, + }; + const suggestion = trino.getSuggestionAtCaretPosition( + commentOtherLine(tokenSql, pos.lineNumber), + pos + )?.keywords; + expect(suggestion).toContain('NOT'); + expect(suggestion).toContain('NOT EXISTS'); + }); }); From 46035498d4732538e5690020af49bc25e66c9e33 Mon Sep 17 00:00:00 2001 From: JackWang032 <2522134117@qq.com> Date: Thu, 20 Mar 2025 14:56:56 +0800 Subject: [PATCH 2/2] chore: add watch script --- package.json | 1 + 1 file changed, 1 insertion(+) diff --git a/package.json b/package.json index 1d010370..66dcaf8b 100644 --- a/package.json +++ b/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",