feat(flink): add flink expression column

DTStack · LuckyFBB · Oct 12, 2024 · Oct 12, 2024 · Oct 16, 2024 · Oct 16, 2024
commit 82a439ea829c23c62dae378902298e0e4b61350f
diff --git a/src/grammar/flink/FlinkSqlParser.g4 b/src/grammar/flink/FlinkSqlParser.g4
@@ -189,6 +189,10 @@ columnName
     | {this.shouldMatchEmpty()}?
     ;
 
+columnNamePath
+    : uid
+    ;
+
 columnNameList
     : LR_BRACKET columnName (COMMA columnName)* RR_BRACKET
     ;
@@ -451,9 +455,9 @@ queryStatement
     : valuesCaluse
     | withClause queryStatement
     | LR_BRACKET queryStatement RR_BRACKET
-    | left=queryStatement operator=(KW_INTERSECT | KW_UNION | KW_EXCEPT) KW_ALL? right=queryStatement orderByCaluse? limitClause?
-    | selectClause orderByCaluse? limitClause?
-    | selectStatement orderByCaluse? limitClause?
+    | left=queryStatement operator=(KW_INTERSECT | KW_UNION | KW_EXCEPT) KW_ALL? right=queryStatement orderByClause? limitClause?
+    | selectClause orderByClause? limitClause?
+    | selectStatement orderByClause? limitClause?
     ;
 
 valuesCaluse
@@ -626,15 +630,15 @@ namedWindow
     ;
 
 windowSpec
-    : name=errorCapturingIdentifier? LR_BRACKET partitionByClause? orderByCaluse? windowFrame? RR_BRACKET
+    : name=errorCapturingIdentifier? LR_BRACKET partitionByClause? orderByClause? windowFrame? RR_BRACKET
     ;
 
 matchRecognizeClause
-    : KW_MATCH_RECOGNIZE LR_BRACKET partitionByClause? orderByCaluse? measuresClause? outputMode? afterMatchStrategy? patternDefination?
+    : KW_MATCH_RECOGNIZE LR_BRACKET partitionByClause? orderByClause? measuresClause? outputMode? afterMatchStrategy? patternDefination?
         patternVariablesDefination RR_BRACKET (KW_AS? identifier)?
     ;
 
-orderByCaluse
+orderByClause
     : KW_ORDER KW_BY orderItemDefition (COMMA orderItemDefition)*
     ;
 
@@ -763,7 +767,7 @@ primaryExpression
     // | identifier '->' expression                                                               #lambda
     // | '(' identifier (',' identifier)+ ')' '->' expression                                     #lambda
     | value=primaryExpression LS_BRACKET index=valueExpression RS_BRACKET # subscript
-    | identifier                                                          # columnReference
+    | columnNamePath                                                      # columnReference
     | dereferenceDefinition                                               # dereference
     | LR_BRACKET expression RR_BRACKET                                    # parenthesizedExpression
     | KW_CURRENT_TIMESTAMP                                                # dateFunctionExpression
@@ -1216,4 +1220,4 @@ nonReservedKeywords
     | KW_WEEK
     | KW_YEARS
     | KW_ZONE
-    ;
+    ;
diff --git a/src/lib/flink/FlinkSqlParser.interp b/src/lib/flink/FlinkSqlParser.interp
diff --git a/src/lib/flink/FlinkSqlParser.ts b/src/lib/flink/FlinkSqlParser.ts
diff --git a/src/lib/flink/FlinkSqlParserListener.ts b/src/lib/flink/FlinkSqlParserListener.ts
@@ -35,6 +35,7 @@ import { ColumnOptionDefinitionContext } from "./FlinkSqlParser.js";
 import { PhysicalColumnDefinitionContext } from "./FlinkSqlParser.js";
 import { ColumnNameCreateContext } from "./FlinkSqlParser.js";
 import { ColumnNameContext } from "./FlinkSqlParser.js";
+import { ColumnNamePathContext } from "./FlinkSqlParser.js";
 import { ColumnNameListContext } from "./FlinkSqlParser.js";
 import { ColumnTypeContext } from "./FlinkSqlParser.js";
 import { LengthOneDimensionContext } from "./FlinkSqlParser.js";
@@ -123,7 +124,7 @@ import { WindowClauseContext } from "./FlinkSqlParser.js";
 import { NamedWindowContext } from "./FlinkSqlParser.js";
 import { WindowSpecContext } from "./FlinkSqlParser.js";
 import { MatchRecognizeClauseContext } from "./FlinkSqlParser.js";
-import { OrderByCaluseContext } from "./FlinkSqlParser.js";
+import { OrderByClauseContext } from "./FlinkSqlParser.js";
 import { OrderItemDefitionContext } from "./FlinkSqlParser.js";
 import { LimitClauseContext } from "./FlinkSqlParser.js";
 import { PartitionByClauseContext } from "./FlinkSqlParser.js";
@@ -497,6 +498,16 @@ export class FlinkSqlParserListener implements ParseTreeListener {
      * @param ctx the parse tree
      */
     exitColumnName?: (ctx: ColumnNameContext) => void;
+    /**
+     * Enter a parse tree produced by `FlinkSqlParser.columnNamePath`.
+     * @param ctx the parse tree
+     */
+    enterColumnNamePath?: (ctx: ColumnNamePathContext) => void;
+    /**
+     * Exit a parse tree produced by `FlinkSqlParser.columnNamePath`.
+     * @param ctx the parse tree
+     */
+    exitColumnNamePath?: (ctx: ColumnNamePathContext) => void;
     /**
      * Enter a parse tree produced by `FlinkSqlParser.columnNameList`.
      * @param ctx the parse tree
@@ -1382,15 +1393,15 @@ export class FlinkSqlParserListener implements ParseTreeListener {
      */
     exitMatchRecognizeClause?: (ctx: MatchRecognizeClauseContext) => void;
     /**
-     * Enter a parse tree produced by `FlinkSqlParser.orderByCaluse`.
+     * Enter a parse tree produced by `FlinkSqlParser.orderByClause`.
      * @param ctx the parse tree
      */
-    enterOrderByCaluse?: (ctx: OrderByCaluseContext) => void;
+    enterOrderByClause?: (ctx: OrderByClauseContext) => void;
     /**
-     * Exit a parse tree produced by `FlinkSqlParser.orderByCaluse`.
+     * Exit a parse tree produced by `FlinkSqlParser.orderByClause`.
      * @param ctx the parse tree
      */
-    exitOrderByCaluse?: (ctx: OrderByCaluseContext) => void;
+    exitOrderByClause?: (ctx: OrderByClauseContext) => void;
     /**
      * Enter a parse tree produced by `FlinkSqlParser.orderItemDefition`.
      * @param ctx the parse tree

diff --git a/src/lib/flink/FlinkSqlParserVisitor.ts b/src/lib/flink/FlinkSqlParserVisitor.ts
@@ -35,6 +35,7 @@ import { ColumnOptionDefinitionContext } from "./FlinkSqlParser.js";
 import { PhysicalColumnDefinitionContext } from "./FlinkSqlParser.js";
 import { ColumnNameCreateContext } from "./FlinkSqlParser.js";
 import { ColumnNameContext } from "./FlinkSqlParser.js";
+import { ColumnNamePathContext } from "./FlinkSqlParser.js";
 import { ColumnNameListContext } from "./FlinkSqlParser.js";
 import { ColumnTypeContext } from "./FlinkSqlParser.js";
 import { LengthOneDimensionContext } from "./FlinkSqlParser.js";
@@ -123,7 +124,7 @@ import { WindowClauseContext } from "./FlinkSqlParser.js";
 import { NamedWindowContext } from "./FlinkSqlParser.js";
 import { WindowSpecContext } from "./FlinkSqlParser.js";
 import { MatchRecognizeClauseContext } from "./FlinkSqlParser.js";
-import { OrderByCaluseContext } from "./FlinkSqlParser.js";
+import { OrderByClauseContext } from "./FlinkSqlParser.js";
 import { OrderItemDefitionContext } from "./FlinkSqlParser.js";
 import { LimitClauseContext } from "./FlinkSqlParser.js";
 import { PartitionByClauseContext } from "./FlinkSqlParser.js";
@@ -392,6 +393,12 @@ export class FlinkSqlParserVisitor<Result> extends AbstractParseTreeVisitor<Resu
      * @return the visitor result
      */
     visitColumnName?: (ctx: ColumnNameContext) => Result;
+    /**
+     * Visit a parse tree produced by `FlinkSqlParser.columnNamePath`.
+     * @param ctx the parse tree
+     * @return the visitor result
+     */
+    visitColumnNamePath?: (ctx: ColumnNamePathContext) => Result;
     /**
      * Visit a parse tree produced by `FlinkSqlParser.columnNameList`.
      * @param ctx the parse tree
@@ -923,11 +930,11 @@ export class FlinkSqlParserVisitor<Result> extends AbstractParseTreeVisitor<Resu
      */
     visitMatchRecognizeClause?: (ctx: MatchRecognizeClauseContext) => Result;
     /**
-     * Visit a parse tree produced by `FlinkSqlParser.orderByCaluse`.
+     * Visit a parse tree produced by `FlinkSqlParser.orderByClause`.
      * @param ctx the parse tree
      * @return the visitor result
      */
-    visitOrderByCaluse?: (ctx: OrderByCaluseContext) => Result;
+    visitOrderByClause?: (ctx: OrderByClauseContext) => Result;
     /**
      * Visit a parse tree produced by `FlinkSqlParser.orderItemDefition`.
      * @param ctx the parse tree

diff --git a/src/parser/flink/flinkErrorListener.ts b/src/parser/flink/flinkErrorListener.ts
@@ -16,6 +16,7 @@ export class FlinkErrorListener extends ParseErrorListener {
         [FlinkSqlParser.RULE_functionName, 'function'],
         [FlinkSqlParser.RULE_functionNameCreate, 'function'],
         [FlinkSqlParser.RULE_columnName, 'column'],
+        [FlinkSqlParser.RULE_columnNamePath, 'column'],
         [FlinkSqlParser.RULE_columnNameCreate, 'column'],
     ]);
 
@@ -60,6 +61,7 @@ export class FlinkErrorListener extends ParseErrorListener {
                     case FlinkSqlParser.RULE_viewPath:
                     case FlinkSqlParser.RULE_functionName:
                     case FlinkSqlParser.RULE_columnName:
+                    case FlinkSqlParser.RULE_columnNamePath:
                     case FlinkSqlParser.RULE_catalogPath: {
                         result.push(`{existing}${name}`);
                         break;

diff --git a/src/parser/flink/index.ts b/src/parser/flink/index.ts
@@ -32,6 +32,7 @@ export class FlinkSQL extends BasicSQL<FlinkSqlLexer, ProgramContext, FlinkSqlPa
         FlinkSqlParser.RULE_functionName, // functionName
         FlinkSqlParser.RULE_functionNameCreate, // functionName that will be created
         FlinkSqlParser.RULE_columnName,
+        FlinkSqlParser.RULE_columnNamePath,
         FlinkSqlParser.RULE_columnNameCreate,
     ]);
 
@@ -111,6 +112,19 @@ export class FlinkSQL extends BasicSQL<FlinkSqlLexer, ProgramContext, FlinkSqlPa
                     syntaxContextType = EntityContextType.COLUMN_CREATE;
                     break;
                 }
+                case FlinkSqlParser.RULE_columnNamePath: {
+                    if (
+                        candidateRule.ruleList.includes(FlinkSqlParser.RULE_selectClause) ||
+                        candidateRule.ruleList.includes(FlinkSqlParser.RULE_whereClause) ||
+                        candidateRule.ruleList.includes(FlinkSqlParser.RULE_groupByClause) ||
+                        candidateRule.ruleList.includes(FlinkSqlParser.RULE_limitClause) ||
+                        candidateRule.ruleList.includes(FlinkSqlParser.RULE_whenClause) ||
+                        candidateRule.ruleList.includes(FlinkSqlParser.RULE_havingClause)
+                    ) {
+                        syntaxContextType = EntityContextType.COLUMN;
+                    }
+                    break;
+                }
                 default:
                     break;
             }

diff --git a/test/parser/flink/errorListener.test.ts b/test/parser/flink/errorListener.test.ts
@@ -7,6 +7,7 @@ const sql3 = `DROP VIEW IF EXIsST aaa aaa`;
 const sql4 = `SELECT * froma aaa`;
 const sql5 = `CREATE VIEW `;
 const sql6 = `DROP CATALOG `;
+const sql7 = `SELECT SUM(amount) FROM Orders GROUP BY length(users) HAVING SUM( `;
 
 describe('FlinkSQL validate invalid sql and test msg', () => {
     const flink = new FlinkSQL();
@@ -64,6 +65,14 @@ describe('FlinkSQL validate invalid sql and test msg', () => {
         );
     });
 
+    test('validate unComplete sql7', () => {
+        const errors = flink.validate(sql7);
+        expect(errors.length).toBe(1);
+        expect(errors[0].message).toBe(
+            `Statement is incomplete, expecting an existing function or an existing column or a keyword`
+        );
+    });
+
     test('validate random text cn', () => {
         flink.locale = 'zh_CN';
         const errors = flink.validate(randomText);
@@ -97,4 +106,12 @@ describe('FlinkSQL validate invalid sql and test msg', () => {
         expect(errors.length).toBe(1);
         expect(errors[0].message).toBe(`'aaa' 在此位置无效，期望一个存在的column或者一个关键字`);
     });
+
+    test('validate unComplete sql7 cn', () => {
+        const errors = flink.validate(sql7);
+        expect(errors.length).toBe(1);
+        expect(errors[0].message).toBe(
+            `语句不完整，期望一个存在的function或者一个存在的column或者一个关键字`
+        );
+    });
 });
diff --git a/test/parser/flink/suggestion/fixtures/syntaxSuggestion.sql b/test/parser/flink/suggestion/fixtures/syntaxSuggestion.sql
@@ -38,4 +38,10 @@ INSERT INTO tb (col, tb.c );
 
 CREATE TABLE yourTable (ts TIMESTAMP(3), WATERMARK FOR  );
 
-CREATE TABLE newTable (  );
+CREATE TABLE newTable (  );
+
+SELECT SUM(amount) FROM Orders GROUP BY length(users) HAVING SUM(amount) > 50;
+
+SELECT * FROM Orders ORDER BY orderTime LIMIT length(order_id);
+
+SELECT age CASE WHEN age < 18 THEN 1  ELSE 0 END AS is_minor FROM dt_catalog.dt_db.subscriptions;
diff --git a/test/parser/flink/suggestion/syntaxSuggestion.test.ts b/test/parser/flink/suggestion/syntaxSuggestion.test.ts
@@ -374,4 +374,89 @@ describe('Flink SQL Syntax Suggestion', () => {
         expect(suggestion).not.toBeUndefined();
         expect(suggestion?.wordRanges.map((token) => token.text)).toEqual([]);
     });
+
+    test('Select expression column', () => {
+        const pos: CaretPosition = {
+            lineNumber: 43,
+            column: 18,
+        };
+        const syntaxes = flink.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === EntityContextType.COLUMN
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['amount']);
+    });
+
+    test('Group by expression column', () => {
+        const pos: CaretPosition = {
+            lineNumber: 43,
+            column: 53,
+        };
+        const syntaxes = flink.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === EntityContextType.COLUMN
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['users']);
+    });
+
+    test('Having expression column', () => {
+        const pos: CaretPosition = {
+            lineNumber: 43,
+            column: 72,
+        };
+        const syntaxes = flink.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === EntityContextType.COLUMN
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['amount']);
+    });
+
+    test('Limit by expression column', () => {
+        const pos: CaretPosition = {
+            lineNumber: 45,
+            column: 62,
+        };
+        const syntaxes = flink.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === EntityContextType.COLUMN
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['order_id']);
+    });
+
+    test('When by expression column', () => {
+        const pos: CaretPosition = {
+            lineNumber: 47,
+            column: 25,
+        };
+        const syntaxes = flink.getSuggestionAtCaretPosition(
+            commentOtherLine(syntaxSql, pos.lineNumber),
+            pos
+        )?.syntax;
+        const suggestion = syntaxes?.find(
+            (syn) => syn.syntaxContextType === EntityContextType.COLUMN
+        );
+
+        expect(suggestion).not.toBeUndefined();
+        expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['age']);
+    });
 });