Skip to content

Commit 9125e8b

Browse files
JackWang032JackWang032
committed
feat: provide follow keywords when get suggestions
1 parent fb50d1a commit 9125e8b

File tree

22 files changed

+341
-88
lines changed

22 files changed

+341
-88
lines changed

‎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+
;

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,4 +67,30 @@ describe('Flink SQL Token Suggestion', () => {
6767
'JARS',
6868
]);
6969
});
70+
71+
test('After CREATE TABLE, show combined keywords', () => {
72+
const pos: CaretPosition = {
73+
lineNumber: 9,
74+
column: 14,
75+
};
76+
const suggestion = flink.getSuggestionAtCaretPosition(
77+
commentOtherLine(tokenSql, pos.lineNumber),
78+
pos
79+
)?.keywords;
80+
expect(suggestion).toContain('IF');
81+
expect(suggestion).toContain('IF NOT EXISTS');
82+
});
83+
84+
test('After CREATE TABLE IF, show combined keywords', () => {
85+
const pos: CaretPosition = {
86+
lineNumber: 9,
87+
column: 17,
88+
};
89+
const suggestion = flink.getSuggestionAtCaretPosition(
90+
commentOtherLine(tokenSql, pos.lineNumber),
91+
pos
92+
)?.keywords;
93+
expect(suggestion).toContain('NOT');
94+
expect(suggestion).toContain('NOT EXISTS');
95+
});
7096
});

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,3 +18,5 @@ LOAD
1818
;
1919
SHOW
2020
;
21+
CREATE TABLE IF NOT EXISTS
22+
;

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@ describe('Hive SQL Token Suggestion', () => {
3333
'MATERIALIZED',
3434
'VIEW',
3535
'TABLE',
36+
'RESOURCE PLAN',
37+
'SCHEDULED QUERY',
38+
'MATERIALIZED VIEW',
3639
]);
3740
});
3841

@@ -68,6 +71,11 @@ describe('Hive SQL Token Suggestion', () => {
6871
'REMOTE',
6972
'DATABASE',
7073
'SCHEMA',
74+
'RESOURCE PLAN',
75+
'SCHEDULED QUERY',
76+
'MATERIALIZED VIEW',
77+
'OR REPLACE',
78+
'MANAGED TABLE',
7179
]);
7280
});
7381

@@ -129,6 +137,9 @@ describe('Hive SQL Token Suggestion', () => {
129137
'TABLE',
130138
'DATABASE',
131139
'SCHEMA',
140+
'RESOURCE PLAN',
141+
'MATERIALIZED VIEW',
142+
'SCHEDULED QUERY',
132143
]);
133144
});
134145

@@ -217,6 +228,36 @@ describe('Hive SQL Token Suggestion', () => {
217228
'EXTENDED',
218229
'DATABASES',
219230
'SCHEMAS',
231+
'CURRENT ROLES',
232+
'ROLE GRANT',
233+
'TABLE EXTENDED',
234+
'MATERIALIZED VIEWS',
220235
]);
221236
});
237+
238+
test('After CREATE TABLE, show combined keywords', () => {
239+
const pos: CaretPosition = {
240+
lineNumber: 21,
241+
column: 14,
242+
};
243+
const suggestion = hive.getSuggestionAtCaretPosition(
244+
commentOtherLine(tokenSql, pos.lineNumber),
245+
pos
246+
)?.keywords;
247+
expect(suggestion).toContain('IF');
248+
expect(suggestion).toContain('IF NOT EXISTS');
249+
});
250+
251+
test('After CREATE TABLE IF, show combined keywords', () => {
252+
const pos: CaretPosition = {
253+
lineNumber: 21,
254+
column: 17,
255+
};
256+
const suggestion = hive.getSuggestionAtCaretPosition(
257+
commentOtherLine(tokenSql, pos.lineNumber),
258+
pos
259+
)?.keywords;
260+
expect(suggestion).toContain('NOT');
261+
expect(suggestion).toContain('NOT EXISTS');
262+
});
222263
});

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,3 +9,5 @@ INSERT ;
99
SHOW ;
1010

1111
CREATE TABLE t1 (id );
12+
13+
CREATE TABLE IF NOT EXISTS;

0 commit comments

Comments
 (0)