Feat/trino sql (#188)

* refactor: prefix trino lexer rule name with KW_

* test: add commentOtherLine function

* feat: optimize trino antlr grammar to adapt to c3

* feat: trinosqlParser supports codeCompletion and spliting

* test: trinoSql codeCompletion unit tests

src/parser/trinosql.ts

@@ -1,9 +1,10 @@
 import { Token } from 'antlr4ts';
 import { CandidatesCollection } from 'antlr4-c3';
 import { TrinoSqlLexer } from '../lib/trinosql/TrinoSqlLexer';
-import { TrinoSqlParser, ProgramContext } from '../lib/trinosql/TrinoSqlParser';
+import { TrinoSqlParser, ProgramContext, StatementContext } from '../lib/trinosql/TrinoSqlParser';
+import { TrinoSqlListener } from '../lib/trinosql/TrinoSqlListener';
 import BasicParser from './common/basicParser';
-import { Suggestions } from './common/basic-parser-types';
+import { Suggestions, SyntaxContextType, SyntaxSuggestion } from './common/basic-parser-types';
 
 export default class TrinoSQL extends BasicParser<TrinoSqlLexer, ProgramContext, TrinoSqlParser> {
     protected createLexerFormCharStream(charStreams) {
@@ -17,19 +18,110 @@ export default class TrinoSQL extends BasicParser<TrinoSqlLexer, ProgramContext,
     }
 
     protected get splitListener() {
-        return null as any;
+        return new TrinoSqlSplitListener();
     }
 
-    protected preferredRules: Set<number> = new Set();
+    protected preferredRules: Set<number> = new Set([
+        TrinoSqlParser.RULE_catalogName,
+        TrinoSqlParser.RULE_catalogNameCreate,
+        TrinoSqlParser.RULE_schemaName,
+        TrinoSqlParser.RULE_schemaNameCreate,
+        TrinoSqlParser.RULE_tableName,
+        TrinoSqlParser.RULE_tableNameCreate,
+        TrinoSqlParser.RULE_viewName,
+        TrinoSqlParser.RULE_viewNameCreate,
+        TrinoSqlParser.RULE_functionName,
+    ]);
 
     protected processCandidates(
         candidates: CandidatesCollection,
         allTokens: Token[],
-        caretTokenIndex: number
+        caretTokenIndex: number,
+        tokenIndexOffset: number
     ): Suggestions<Token> {
+        const originalSyntaxSuggestions: SyntaxSuggestion<Token>[] = [];
+        const keywords: string[] = [];
+
+        for (let candidate of candidates.rules) {
+            const [ruleType, candidateRule] = candidate;
+            const startTokenIndex = candidateRule.startTokenIndex + tokenIndexOffset;
+            const tokenRanges = allTokens.slice(
+                startTokenIndex,
+                caretTokenIndex + tokenIndexOffset + 1
+            );
+
+            let syntaxContextType: SyntaxContextType;
+            switch (ruleType) {
+                case TrinoSqlParser.RULE_catalogName: {
+                    syntaxContextType = SyntaxContextType.CATALOG;
+                    break;
+                }
+                case TrinoSqlParser.RULE_schemaName: {
+                    syntaxContextType = SyntaxContextType.DATABASE;
+                    break;
+                }
+                case TrinoSqlParser.RULE_schemaNameCreate: {
+                    syntaxContextType = SyntaxContextType.DATABASE_CREATE;
+                    break;
+                }
+                case TrinoSqlParser.RULE_tableName: {
+                    syntaxContextType = SyntaxContextType.TABLE;
+                    break;
+                }
+                case TrinoSqlParser.RULE_tableNameCreate: {
+                    syntaxContextType = SyntaxContextType.TABLE_CREATE;
+                    break;
+                }
+                case TrinoSqlParser.RULE_viewName: {
+                    syntaxContextType = SyntaxContextType.VIEW;
+                    break;
+                }
+                case TrinoSqlParser.RULE_viewNameCreate: {
+                    syntaxContextType = SyntaxContextType.VIEW_CREATE;
+                    break;
+                }
+                case TrinoSqlParser.RULE_functionName: {
+                    syntaxContextType = SyntaxContextType.FUNCTION;
+                    break;
+                }
+                default:
+                    break;
+            }
+
+            if (syntaxContextType) {
+                originalSyntaxSuggestions.push({
+                    syntaxContextType,
+                    wordRanges: tokenRanges,
+                });
+            }
+        }
+
+        for (let candidate of candidates.tokens) {
+            const symbolicName = this._parser.vocabulary.getSymbolicName(candidate[0]);
+            const displayName = this._parser.vocabulary.getDisplayName(candidate[0]);
+            if (symbolicName && symbolicName.startsWith('KW_')) {
+                const keyword =
+                    displayName.startsWith("'") && displayName.endsWith("'")
+                        ? displayName.slice(1, -1)
+                        : displayName;
+                keywords.push(keyword);
+            }
+        }
         return {
-            syntax: [],
-            keywords: [],
+            syntax: originalSyntaxSuggestions,
+            keywords,
         };
     }
 }
+
+export class TrinoSqlSplitListener implements TrinoSqlListener {
+    private _statementsContext: StatementContext[] = [];
+
+    exitStatement = (ctx: StatementContext) => {
+        this._statementsContext.push(ctx);
+    };
+
+    get statementsContext() {
+        return this._statementsContext;
+    }
+}

test/helper.ts

@@ -68,3 +68,15 @@ export function getReportTableHeader(title: string) {
 export function exportReportTable(markdown: string, output: string) {
     fs.writeFileSync(path.join(output, 'benchmark.md'), markdown);
 }
+
+export function commentOtherLine(sqlContent: string, line: number) {
+    const slices = sqlContent.split('\n').map((item, index) => {
+        if (index !== line - 1) {
+            return '-- ' + item;
+        } else {
+            return item;
+        }
+    });
+
+    return slices.join('\n');
+}

test/parser/flinksql/suggestion/syntaxSuggestion.test.ts

@@ -55,7 +55,7 @@ describe('Flink SQL Syntax Suggestion', () => {
         const suggestion = syntaxes?.find(
             (syn) => syn.syntaxContextType === SyntaxContextType.TABLE
         );
-        console.log(syntaxes);
+
         expect(suggestion).not.toBeUndefined();
         expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['cat', '.']);
     });

test/parser/trinosql/suggestion/fixtures/syntaxSuggestion.sql

@@ -0,0 +1,19 @@
+INSERT INTO db.tb ;
+
+SELECT ids FROM db.;
+
+CREATE TABLE db. VALUES;
+
+DROP TABLE IF EXISTS db.a;
+
+CREATE OR REPLACE VIEW db.v;
+
+DROP VIEW db.v ;
+
+SELECT name, calculate_age(birthday) AS age FROM students;
+
+CREATE SCHEMA db ;
+
+DROP SCHEMA IF EXISTS sch;
+
+SHOW COLUMNS FROM tb ;

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

@@ -0,0 +1,13 @@
+ALTER ;
+
+CREATE ;
+
+DEALLOCATE ;
+
+DELETE ;
+
+DESCRIBE ;
+
+DROP ;
+
+INSERT ;

test/parser/trinosql/suggestion/syntaxSuggestion.test.ts

@@ -0,0 +1,201 @@
+import fs from 'fs';
+import path from 'path';
+import { CaretPosition, SyntaxContextType } from '../../../../src/parser/common/basic-parser-types';
+import TrinoSQL from '../../../../src/parser/trinosql';
+import { commentOtherLine } from '../../../helper';
+
+const syntaxSql = fs.readFileSync(
+    path.join(__dirname, 'fixtures', 'syntaxSuggestion.sql'),
+    'utf-8'
+);
+
+describe('Trino SQL Syntax Suggestion', () => {
+    const parser = new TrinoSQL();
+
+    test('Validate Syntax SQL', () => {
+        expect(parser.validate(syntaxSql).length).not.toBe(0);
+        expect(parser.validate(syntaxSql).length).not.toBe(0);
+        expect(parser.validate(syntaxSql).length).not.toBe(0);
+    });
+
+    test('Insert table ', () => {
+        const pos: CaretPosition = {
+            lineNumber: 1,
+            column: 18,
+        };
+        const syntaxes = parser.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === SyntaxContextType.TABLE
+        );
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['db', '.', 'tb']);
+    });
+
+    test('Select table ', () => {
+        const pos: CaretPosition = {
+            lineNumber: 3,
+            column: 20,
+        };
+        const syntaxes =
+            parser.getSuggestionAtCaretPosition(commentOtherLine(syntaxSql, pos.lineNumber), pos)
+                ?.syntax ?? [];
+
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === SyntaxContextType.TABLE
+        );
+
+        expect(
+            syntaxes.some((item) => item.syntaxContextType === SyntaxContextType.VIEW)
+        ).toBeTruthy();
+        expect(
+            syntaxes.some((item) => item.syntaxContextType === SyntaxContextType.FUNCTION)
+        ).toBeTruthy();
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['db', '.']);
+    });
+
+    test('Create table ', () => {
+        const pos: CaretPosition = {
+            lineNumber: 5,
+            column: 17,
+        };
+        const syntaxes = parser.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === SyntaxContextType.TABLE_CREATE
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['db', '.']);
+    });
+
+    test('DROP table ', () => {
+        const pos: CaretPosition = {
+            lineNumber: 7,
+            column: 26,
+        };
+        const syntaxes = parser.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === SyntaxContextType.TABLE
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['db', '.', 'a']);
+    });
+
+    test('Create view ', () => {
+        const pos: CaretPosition = {
+            lineNumber: 9,
+            column: 28,
+        };
+        const syntaxes = parser.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === SyntaxContextType.VIEW_CREATE
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['db', '.', 'v']);
+    });
+
+    test('Drop view ', () => {
+        const pos: CaretPosition = {
+            lineNumber: 11,
+            column: 15,
+        };
+        const syntaxes = parser.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === SyntaxContextType.VIEW
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['db', '.', 'v']);
+    });
+
+    test('Use function', () => {
+        const pos: CaretPosition = {
+            lineNumber: 13,
+            column: 27,
+        };
+        const syntaxes = parser.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === SyntaxContextType.FUNCTION
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['calculate_age']);
+    });
+
+    test('Create schema', () => {
+        const pos: CaretPosition = {
+            lineNumber: 15,
+            column: 17,
+        };
+        const syntaxes = parser.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === SyntaxContextType.DATABASE_CREATE
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['db']);
+    });
+
+    test('Drop schema', () => {
+        const pos: CaretPosition = {
+            lineNumber: 17,
+            column: 26,
+        };
+        const syntaxes = parser.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === SyntaxContextType.DATABASE
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['sch']);
+    });
+
+    test('Show Column From', () => {
+        const pos: CaretPosition = {
+            lineNumber: 19,
+            column: 21,
+        };
+        const syntaxes =
+            parser.getSuggestionAtCaretPosition(commentOtherLine(syntaxSql, pos.lineNumber), pos)
+                ?.syntax ?? [];
+
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === SyntaxContextType.TABLE
+        );
+        expect(
+            syntaxes.some((item) => item.syntaxContextType === SyntaxContextType.VIEW)
+        ).toBeTruthy();
+        expect(
+            syntaxes.some((item) => item.syntaxContextType === SyntaxContextType.FUNCTION)
+        ).toBeTruthy();
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['tb']);
+    });
+});