Skip to content

Commit bcf8369

Browse files
JackWang032JackWang032
authored
Feat/follow keywords (#407)
* feat: provide follow keywords when get suggestions * chore: add watch script
1 parent fb50d1a commit bcf8369

File tree

23 files changed

+342
-88
lines changed

23 files changed

+342
-88
lines changed

‎package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
"prepublishOnly": "npm run build",
2828
"antlr4": "node ./scripts/antlr4.js",
2929
"build": "rm -rf dist && tsc",
30+
"watch": "tsc -w",
3031
"check-types": "tsc -p ./tsconfig.json && tsc -p ./test/tsconfig.json",
3132
"test": "NODE_OPTIONS=--max_old_space_size=4096 && jest",
3233
"release": "node ./scripts/release.js",

‎src/parser/common/tokenUtils.ts

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
/**
2+
* Utility function for processing SQL tokens and generating keyword suggestions
3+
*/
4+
5+
import { Parser } from 'antlr4ng';
6+
import { CandidatesCollection } from 'antlr4-c3';
7+
8+
/**
9+
* Process token candidates and generate a list of keyword suggestions
10+
* @param parser SQL parser instance
11+
* @param tokens token candidates
12+
* @returns list of keyword suggestions
13+
*/
14+
export function processTokenCandidates(
15+
parser: Parser,
16+
tokens: CandidatesCollection['tokens']
17+
): string[] {
18+
const keywords: string[] = [];
19+
20+
const cleanDisplayName = (displayName: string | null): string => {
21+
return displayName && displayName.startsWith("'") && displayName.endsWith("'")
22+
? displayName.slice(1, -1)
23+
: displayName || '';
24+
};
25+
26+
const isKeywordToken = (token: number): boolean => {
27+
const symbolicName = parser.vocabulary.getSymbolicName(token);
28+
return Boolean(symbolicName?.startsWith('KW_'));
29+
};
30+
31+
for (const [token, followSets] of tokens) {
32+
const displayName = parser.vocabulary.getDisplayName(token);
33+
34+
if (!displayName || !isKeywordToken(token)) continue;
35+
36+
const keyword = cleanDisplayName(displayName);
37+
keywords.push(keyword);
38+
39+
if (followSets.length && followSets.every((s) => isKeywordToken(s))) {
40+
const followKeywords = followSets
41+
.map((s) => cleanDisplayName(parser.vocabulary.getDisplayName(s)))
42+
.filter(Boolean);
43+
44+
if (followKeywords.length) {
45+
const combinedKeyword = [keyword, ...followKeywords].join(' ');
46+
keywords.push(combinedKeyword);
47+
}
48+
}
49+
}
50+
51+
return keywords;
52+
}

‎src/parser/flink/index.ts

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { CandidatesCollection } from 'antlr4-c3';
22
import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
3-
3+
import { processTokenCandidates } from '../common/tokenUtils';
44
import { FlinkSqlLexer } from '../../lib/flink/FlinkSqlLexer';
55
import { FlinkSqlParser, ProgramContext } from '../../lib/flink/FlinkSqlParser';
66
import { BasicSQL } from '../common/basicSQL';
@@ -123,17 +123,9 @@ export class FlinkSQL extends BasicSQL<FlinkSqlLexer, ProgramContext, FlinkSqlPa
123123
}
124124
}
125125

126-
for (let candidate of candidates.tokens) {
127-
const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
128-
const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
129-
if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
130-
const keyword =
131-
displayName.startsWith("'") && displayName.endsWith("'")
132-
? displayName.slice(1, -1)
133-
: displayName;
134-
keywords.push(keyword);
135-
}
136-
}
126+
const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
127+
keywords.push(...processedKeywords);
128+
137129
return {
138130
syntax: originalSyntaxSuggestions,
139131
keywords,

‎src/parser/hive/index.ts

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { CandidatesCollection } from 'antlr4-c3';
22
import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
3-
3+
import { processTokenCandidates } from '../common/tokenUtils';
44
import { HiveSqlLexer } from '../../lib/hive/HiveSqlLexer';
55
import { HiveSqlParser, ProgramContext } from '../../lib/hive/HiveSqlParser';
66
import { BasicSQL } from '../common/basicSQL';
@@ -119,17 +119,9 @@ export class HiveSQL extends BasicSQL<HiveSqlLexer, ProgramContext, HiveSqlParse
119119
}
120120
}
121121

122-
for (let candidate of candidates.tokens) {
123-
const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
124-
const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
125-
if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
126-
const keyword =
127-
displayName.startsWith("'") && displayName.endsWith("'")
128-
? displayName.slice(1, -1)
129-
: displayName;
130-
keywords.push(keyword);
131-
}
132-
}
122+
const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
123+
keywords.push(...processedKeywords);
124+
133125
return {
134126
syntax: originalSyntaxSuggestions,
135127
keywords,

‎src/parser/impala/index.ts

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { CandidatesCollection } from 'antlr4-c3';
22
import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
3-
3+
import { processTokenCandidates } from '../common/tokenUtils';
44
import { ImpalaSqlLexer } from '../../lib/impala/ImpalaSqlLexer';
55
import { ImpalaSqlParser, ProgramContext } from '../../lib/impala/ImpalaSqlParser';
66
import { BasicSQL } from '../common/basicSQL';
@@ -116,17 +116,9 @@ export class ImpalaSQL extends BasicSQL<ImpalaSqlLexer, ProgramContext, ImpalaSq
116116
}
117117
}
118118

119-
for (let candidate of candidates.tokens) {
120-
const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
121-
const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
122-
if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
123-
const keyword =
124-
displayName.startsWith("'") && displayName.endsWith("'")
125-
? displayName.slice(1, -1)
126-
: displayName;
127-
keywords.push(keyword);
128-
}
129-
}
119+
const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
120+
keywords.push(...processedKeywords);
121+
130122
return {
131123
syntax: originalSyntaxSuggestions,
132124
keywords,

‎src/parser/mysql/index.ts

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { CandidatesCollection } from 'antlr4-c3';
22
import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
3-
3+
import { processTokenCandidates } from '../common/tokenUtils';
44
import { MySqlLexer } from '../../lib/mysql/MySqlLexer';
55
import { MySqlParser, ProgramContext } from '../../lib/mysql/MySqlParser';
66
import { BasicSQL } from '../common/basicSQL';
@@ -118,17 +118,8 @@ export class MySQL extends BasicSQL<MySqlLexer, ProgramContext, MySqlParser> {
118118
}
119119
}
120120

121-
for (const candidate of candidates.tokens) {
122-
const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
123-
const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
124-
if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
125-
const keyword =
126-
displayName.startsWith("'") && displayName.endsWith("'")
127-
? displayName.slice(1, -1)
128-
: displayName;
129-
keywords.push(keyword);
130-
}
131-
}
121+
const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
122+
keywords.push(...processedKeywords);
132123

133124
return {
134125
syntax: originalSyntaxSuggestions,

‎src/parser/postgresql/index.ts

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { CandidatesCollection } from 'antlr4-c3';
22
import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
3+
import { processTokenCandidates } from '../common/tokenUtils';
34

45
import { PostgreSqlLexer } from '../../lib/postgresql/PostgreSqlLexer';
56
import { PostgreSqlParser, ProgramContext } from '../../lib/postgresql/PostgreSqlParser';
@@ -137,17 +138,9 @@ export class PostgreSQL extends BasicSQL<PostgreSqlLexer, ProgramContext, Postgr
137138
}
138139
}
139140

140-
for (let candidate of candidates.tokens) {
141-
const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
142-
const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
143-
if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
144-
const keyword =
145-
displayName.startsWith("'") && displayName.endsWith("'")
146-
? displayName.slice(1, -1)
147-
: displayName;
148-
keywords.push(keyword);
149-
}
150-
}
141+
const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
142+
keywords.push(...processedKeywords);
143+
151144
return {
152145
syntax: originalSyntaxSuggestions,
153146
keywords,

‎src/parser/spark/index.ts

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { CandidatesCollection } from 'antlr4-c3';
22
import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
3-
3+
import { processTokenCandidates } from '../common/tokenUtils';
44
import { SparkSqlLexer } from '../../lib/spark/SparkSqlLexer';
55
import { ProgramContext, SparkSqlParser } from '../../lib/spark/SparkSqlParser';
66
import { BasicSQL } from '../common/basicSQL';
@@ -118,17 +118,8 @@ export class SparkSQL extends BasicSQL<SparkSqlLexer, ProgramContext, SparkSqlPa
118118
}
119119
}
120120

121-
for (const candidate of candidates.tokens) {
122-
const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
123-
const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
124-
if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
125-
const keyword =
126-
displayName.startsWith("'") && displayName.endsWith("'")
127-
? displayName.slice(1, -1)
128-
: displayName;
129-
keywords.push(keyword);
130-
}
131-
}
121+
const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
122+
keywords.push(...processedKeywords);
132123

133124
return {
134125
syntax: originalSyntaxSuggestions,

‎src/parser/trino/index.ts

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { CandidatesCollection } from 'antlr4-c3';
22
import { CharStream, CommonTokenStream, Token } from 'antlr4ng';
3-
3+
import { processTokenCandidates } from '../common/tokenUtils';
44
import { TrinoSqlLexer } from '../../lib/trino/TrinoSqlLexer';
55
import { ProgramContext, TrinoSqlParser } from '../../lib/trino/TrinoSqlParser';
66
import { BasicSQL } from '../common/basicSQL';
@@ -128,17 +128,9 @@ export class TrinoSQL extends BasicSQL<TrinoSqlLexer, ProgramContext, TrinoSqlPa
128128
}
129129
}
130130

131-
for (let candidate of candidates.tokens) {
132-
const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
133-
const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
134-
if (displayName && symbolicName && symbolicName.startsWith('KW_')) {
135-
const keyword =
136-
displayName.startsWith("'") && displayName.endsWith("'")
137-
? displayName.slice(1, -1)
138-
: displayName;
139-
keywords.push(keyword);
140-
}
141-
}
131+
const processedKeywords = processTokenCandidates(this._parser, candidates.tokens);
132+
keywords.push(...processedKeywords);
133+
142134
return {
143135
syntax: originalSyntaxSuggestions,
144136
keywords,

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,4 +4,7 @@ USE
44
;
55
CREATE
66
;
7-
SHOW
7+
SHOW
8+
;
9+
CREATE TABLE IF NOT EXISTS
10+
;

0 commit comments

Comments
 (0)