Skip to content

Commit 435166d

Browse files
committed
Improved support for parser conditional and additional statement types
Signed-off-by: worksofliam <[email protected]>
1 parent 00428ed commit 435166d

File tree

6 files changed

+391
-185
lines changed

6 files changed

+391
-185
lines changed

src/language/sql/document.ts

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ export default class Document {
3535
let statementStart = 0;
3636

3737
for (let i = 0; i < tokens.length; i++) {
38+
const upperValue = tokens[i].value?.toUpperCase();
3839
switch (tokens[i].type) {
3940
case `semicolon`:
4041
const statementTokens = tokens.slice(statementStart, i);
@@ -45,19 +46,25 @@ export default class Document {
4546
break;
4647

4748
case `statementType`:
48-
currentStatementType = StatementTypeWord[tokens[i].value?.toUpperCase()];
49+
currentStatementType = StatementTypeWord[upperValue];
4950
break;
5051

5152
case `keyword`:
52-
switch (tokens[i].value?.toUpperCase()) {
53+
switch (upperValue) {
5354
case `LOOP`:
54-
// This handles the case that 'END LOOP' is supported.
55-
if (currentStatementType === StatementType.End) {
56-
break;
57-
}
55+
case `THEN`:
5856
case `BEGIN`:
5957
case `DO`:
60-
case `THEN`:
58+
// This handles the case that 'END LOOP' is supported.
59+
if (upperValue === `LOOP` && currentStatementType === StatementType.End) {
60+
break;
61+
}
62+
63+
// Support for THEN in conditionals
64+
if (upperValue === `THEN` && !Statement.typeIsConditional(currentStatementType)) {
65+
break;
66+
}
67+
6168
// We include BEGIN in the current statement
6269
// then the next statement beings
6370
const statementTokens = tokens.slice(statementStart, i+1);
@@ -102,7 +109,7 @@ export default class Document {
102109
let depth = 0;
103110

104111
for (const statement of this.statements) {
105-
if (statement.isBlockEnder()) {
112+
if (statement.isCompoundEnd()) {
106113
if (depth > 0) {
107114
currentGroup.push(statement);
108115

@@ -118,7 +125,7 @@ export default class Document {
118125
currentGroup = [];
119126
}
120127
} else
121-
if (statement.isBlockOpener()) {
128+
if (statement.isCompoundStart()) {
122129
if (depth > 0) {
123130
currentGroup.push(statement);
124131
} else {

src/language/sql/statement.ts

Lines changed: 58 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ const tokenIs = (token: Token|undefined, type: string, value?: string) => {
77

88
export default class Statement {
99
public type: StatementType = StatementType.Unknown;
10+
private label: string|undefined;
1011

1112
constructor(public tokens: Token[], public range: IRange) {
1213
this.tokens = this.tokens.filter(newToken => newToken.type !== `newline`);
@@ -19,11 +20,15 @@ export default class Statement {
1920
first = this.tokens[2];
2021
}
2122

22-
if (tokenIs(first, `statementType`) || tokenIs(first, `keyword`, `END`) || tokenIs(first, `keyword`, `BEGIN`)) {
23-
const wordValue = first.value?.toUpperCase();
24-
25-
this.type = StatementTypeWord[wordValue];
23+
if (tokenIs(first, `word`) && tokenIs(this.tokens[1], `colon`)) {
24+
// Possible label?
25+
this.label = first.value;
26+
first = this.tokens[2];
2627
}
28+
29+
const wordValue = first.value?.toUpperCase();
30+
31+
this.type = StatementTypeWord[wordValue] || StatementType.Unknown;
2732

2833
switch (this.type) {
2934
case StatementType.Create:
@@ -35,7 +40,7 @@ export default class Statement {
3540
}
3641
}
3742

38-
isBlockOpener() {
43+
isCompoundStart() {
3944
if (this.tokens.length === 1 && tokenIs(this.tokens[0], `keyword`, `BEGIN`)) {
4045
return true;
4146
}
@@ -51,7 +56,19 @@ export default class Statement {
5156
return false;
5257
}
5358

54-
isBlockEnder() {
59+
static typeIsConditional(type: StatementType) {
60+
return [StatementType.If, StatementType.While, StatementType.Loop, StatementType.For].includes(type);
61+
}
62+
63+
isConditionStart() {
64+
return Statement.typeIsConditional(this.type);
65+
}
66+
67+
isConditionEnd() {
68+
return this.type === StatementType.End && this.tokens.length > 1;
69+
}
70+
71+
isCompoundEnd() {
5572
return this.type === StatementType.End && this.tokens.length === 1;
5673
}
5774

@@ -314,17 +331,21 @@ export default class Statement {
314331
}
315332

316333
const basicQueryFinder = (startIndex: number): void => {
334+
let currentClause: undefined|"select"|"from";
317335
for (let i = startIndex; i < this.tokens.length; i++) {
318336
if (tokenIs(this.tokens[i], `clause`, `FROM`)) {
319-
inFromClause = true;
320-
} else if (inFromClause && tokenIs(this.tokens[i], `clause`) || tokenIs(this.tokens[i], `join`) || tokenIs(this.tokens[i], `closebracket`)) {
321-
inFromClause = false;
337+
currentClause = `from`;
338+
}
339+
else if (tokenIs(this.tokens[i], `statementType`, `SELECT`)) {
340+
currentClause = `select`;
341+
} else if (currentClause === `from` && tokenIs(this.tokens[i], `clause`) || tokenIs(this.tokens[i], `join`) || tokenIs(this.tokens[i], `closebracket`)) {
342+
currentClause = undefined;
322343
}
323344

324345
if (tokenIs(this.tokens[i], `clause`, `FROM`) ||
325346
(this.type !== StatementType.Select && tokenIs(this.tokens[i], `clause`, `INTO`)) ||
326347
tokenIs(this.tokens[i], `join`) ||
327-
(inFromClause && tokenIs(this.tokens[i], `comma`)
348+
(currentClause === `from` && tokenIs(this.tokens[i], `comma`)
328349
)) {
329350
const sqlObj = this.getRefAtToken(i+1);
330351
if (sqlObj) {
@@ -334,6 +355,15 @@ export default class Statement {
334355
i += 3; //For the brackets
335356
}
336357
}
358+
} else if (currentClause === `select` && tokenIs(this.tokens[i], `function`)) {
359+
const sqlObj = this.getRefAtToken(i);
360+
if (sqlObj) {
361+
doAdd(sqlObj);
362+
i += sqlObj.tokens.length;
363+
if (sqlObj.isUDTF || sqlObj.fromLateral) {
364+
i += 3; //For the brackets
365+
}
366+
}
337367
}
338368
}
339369
}
@@ -592,7 +622,7 @@ export default class Statement {
592622
}
593623

594624
if (options.withSystemName) {
595-
if (tokenIs(this.tokens[endIndex+1], `keyword`, `FOR`) && tokenIs(this.tokens[endIndex+2], `word`, `SYSTEM`) && tokenIs(this.tokens[endIndex+3], `word`, `NAME`)) {
625+
if (tokenIs(this.tokens[endIndex+1], `statementType`, `FOR`) && tokenIs(this.tokens[endIndex+2], `word`, `SYSTEM`) && tokenIs(this.tokens[endIndex+3], `word`, `NAME`)) {
596626
if (this.tokens[endIndex+4] && NameTypes.includes(this.tokens[endIndex+4].type)) {
597627
sqlObj.object.system = this.tokens[endIndex+4].value;
598628
}
@@ -624,10 +654,25 @@ export default class Statement {
624654

625655
switch (currentToken.type) {
626656
case `statementType`:
627-
if (declareStmt) continue;
657+
const currentValue = currentToken.value.toLowerCase();
658+
if (declareStmt) {
659+
if (currentValue === `for`) {
660+
ranges.push({
661+
type: `remove`,
662+
range: {
663+
start: declareStmt.range.start,
664+
end: currentToken.range.end
665+
}
666+
});
667+
668+
declareStmt = undefined;
669+
}
670+
671+
continue;
672+
};
628673

629674
// If we're in a DECLARE, it's likely a cursor definition
630-
if (currentToken.value.toLowerCase() === `declare`) {
675+
if (currentValue === `declare`) {
631676
declareStmt = currentToken;
632677
}
633678
break;
@@ -716,19 +761,6 @@ export default class Statement {
716761
}
717762
});
718763
}
719-
} else
720-
if (declareStmt && tokenIs(currentToken, `keyword`, `FOR`)) {
721-
// If we're a DECLARE, and we found the FOR keyword, the next
722-
// set of tokens should be the select.
723-
ranges.push({
724-
type: `remove`,
725-
range: {
726-
start: declareStmt.range.start,
727-
end: currentToken.range.end
728-
}
729-
});
730-
731-
declareStmt = undefined;
732764
}
733765
break;
734766
}

src/language/sql/tests/blocks.test.ts

Lines changed: 58 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11

2-
import { assert, describe, expect, test } from 'vitest'
2+
import { describe, expect, test } from 'vitest'
33
import Document from '../document';
4-
import { StatementType } from '../types';
54

6-
describe(`Block statement tests`, () => {
5+
const parserScenarios = describe.each([
6+
{newDoc: (content: string) => new Document(content), isFormatted: false},
7+
]);
8+
9+
parserScenarios(`Block statement tests`, ({newDoc, isFormatted}) => {
710
test('Block start tests', () => {
811
const lines = [
912
`CREATE ALIAS "TestDelimiters"."Delimited Alias" FOR "TestDelimiters"."Delimited Table";`,
@@ -15,19 +18,19 @@ describe(`Block statement tests`, () => {
1518
`LANGUAGE SQL BEGIN SET "Delimited Parameter" = 13; END;`,
1619
].join(`\n`);
1720

18-
const doc = new Document(lines);
21+
const doc = newDoc(lines);
1922

2023
// CREATE, CREATE, RETURN, END, CREATE, SET, END
2124
expect(doc.statements.length).toBe(7);
2225

2326
const aliasDef = doc.statements[0];
24-
expect(aliasDef.isBlockOpener()).toBeFalsy();
27+
expect(aliasDef.isCompoundStart()).toBeFalsy();
2528

2629
const functionDef = doc.statements[1];
27-
expect(functionDef.isBlockOpener()).toBeTruthy();
30+
expect(functionDef.isCompoundStart()).toBeTruthy();
2831

2932
const procedureDef = doc.statements[4];
30-
expect(procedureDef.isBlockOpener()).toBeTruthy();
33+
expect(procedureDef.isCompoundStart()).toBeTruthy();
3134
});
3235

3336
test('Compound statement test', () => {
@@ -53,21 +56,21 @@ describe(`Block statement tests`, () => {
5356
`LANGUAGE SQL BEGIN SET "Delimited Parameter" = 13; END;`,
5457
].join(`\n`);
5558

56-
const doc = new Document(lines);
59+
const doc = newDoc(lines);
5760

5861
const t = doc.statements.length;
5962

6063
const aliasDef = doc.statements[0];
61-
expect(aliasDef.isBlockOpener()).toBeFalsy();
64+
expect(aliasDef.isCompoundStart()).toBeFalsy();
6265

6366
const functionDef = doc.statements[1];
64-
expect(functionDef.isBlockOpener()).toBeTruthy();
67+
expect(functionDef.isCompoundStart()).toBeTruthy();
6568

6669
const functionEnd = doc.statements[3];
67-
expect(functionEnd.isBlockEnder()).toBeTruthy();
70+
expect(functionEnd.isCompoundEnd()).toBeTruthy();
6871

6972
const beginBlock = doc.statements[4];
70-
expect(beginBlock.isBlockOpener()).toBeTruthy();
73+
expect(beginBlock.isCompoundStart()).toBeTruthy();
7174
});
7275

7376
test('Statement groups', () => {
@@ -96,23 +99,58 @@ describe(`Block statement tests`, () => {
9699
`LANGUAGE SQL BEGIN SET "Delimited Parameter" = 13; END;`,
97100
].join(`\r\n`);
98101

99-
const doc = new Document(lines);
102+
const doc = newDoc(lines);
100103

101104
const groups = doc.getStatementGroups();
102105

103106
expect(groups.length).toBe(4);
104107

105108
const aliasStatement = groups[0];
106-
const aliasSubstring = lines.substring(aliasStatement.range.start, aliasStatement.range.end);
109+
const aliasSubstring = doc.content.substring(aliasStatement.range.start, aliasStatement.range.end);
107110
expect(aliasSubstring).toBe(`CREATE ALIAS "TestDelimiters"."Delimited Alias" FOR "TestDelimiters"."Delimited Table"`);
108111

112+
const functionStatement = groups[1];
113+
const functionSubstring = doc.content.substring(functionStatement.range.start, functionStatement.range.end);
114+
115+
if (isFormatted) {
116+
expect(functionSubstring).toBe([
117+
`CREATE FUNCTION "TestDelimiters"."Delimited Function"(`,
118+
` "Delimited Parameter" INTEGER`,
119+
`) RETURNS INTEGER LANGUAGE SQL BEGIN`,
120+
` RETURN "Delimited Parameter";`,
121+
`END`,
122+
].join(`\r\n`));
123+
} else {
124+
expect(functionSubstring).toBe([
125+
`CREATE FUNCTION "TestDelimiters"."Delimited Function" ("Delimited Parameter" INTEGER) `,
126+
`RETURNS INTEGER LANGUAGE SQL BEGIN RETURN "Delimited Parameter"; END`
127+
].join(`\r\n`))
128+
}
109129
const beginStatement = groups[2];
110-
const compoundSubstring = lines.substring(beginStatement.range.start, beginStatement.range.end);
111-
expect(compoundSubstring).toBe(compoundStatement);
130+
expect(beginStatement.statements.length).toBe(9);
131+
const compoundSubstring = doc.content.substring(beginStatement.range.start, beginStatement.range.end);
132+
133+
if (isFormatted) {
134+
expect(compoundSubstring).toBe([
135+
`BEGIN`,
136+
` DECLARE already_exists SMALLINT DEFAULT 0;`,
137+
` DECLARE dup_object_hdlr CONDITION FOR SQLSTATE '42710';`,
138+
` DECLARE CONTINUE HANDLER FOR dup_object_hdlr SET already_exists = 1;`,
139+
` CREATE TABLE table1(`,
140+
` col1 INT`,
141+
` );`,
142+
` IF already_exists > 0 THEN;`,
143+
` DELETE FROM table1;`,
144+
` END IF;`,
145+
`END`,
146+
].join(`\r\n`));
147+
} else {
148+
expect(compoundSubstring).toBe(compoundStatement);
149+
}
112150
});
113151
});
114152

115-
describe(`Definition tests`, () => {
153+
parserScenarios(`Definition tests`, ({newDoc}) => {
116154
test(`Alias, function, procedure`, () => {
117155
const lines = [
118156
`CREATE ALIAS "TestDelimiters"."Delimited Alias" FOR "TestDelimiters"."Delimited Table";`,
@@ -124,7 +162,7 @@ describe(`Definition tests`, () => {
124162
`LANGUAGE SQL BEGIN SET "Delimited Parameter" = 13; END;`,
125163
].join(`\n`);
126164

127-
const doc = new Document(lines);
165+
const doc = newDoc(lines);
128166

129167
const defs = doc.getDefinitions();
130168

@@ -161,7 +199,7 @@ describe(`Definition tests`, () => {
161199
`END;`,
162200
].join(`\r\n`);
163201

164-
const doc = new Document(lines);
202+
const doc = newDoc(lines);
165203

166204
const defs = doc.getDefinitions();
167205

@@ -245,7 +283,7 @@ describe(`Definition tests`, () => {
245283
`END ; `,
246284
].join(`\n`);
247285

248-
const doc = new Document(lines);
286+
const doc = newDoc(lines);
249287

250288
const groups = doc.getStatementGroups();
251289
expect(groups.length).toBe(1);

0 commit comments

Comments
 (0)