From 009effc6868933e416162af6afe5f58cf0e48ed5 Mon Sep 17 00:00:00 2001
From: zhaoge <>
Date: Thu, 4 Jul 2024 17:47:50 +0800
Subject: [PATCH] feat: collect alias in select statement

---
 src/parser/common/entityCollector.ts          |  4 +++
 src/parser/flink/flinkEntityCollector.ts      | 12 ++++++-
 .../contextCollect/entityCollector.test.ts    | 32 ++++++++++++-------
 .../flink/contextCollect/fixtures/common.sql  |  2 +-
 4 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/src/parser/common/entityCollector.ts b/src/parser/common/entityCollector.ts
index 2a8adb0e..77588786 100644
--- a/src/parser/common/entityCollector.ts
+++ b/src/parser/common/entityCollector.ts
@@ -277,6 +277,10 @@ export abstract class EntityCollector {
 
     exitEveryRule() {}
 
+    getRootStmt() {
+        return this._rootStmt;
+    }
+
     getEntities() {
         return Array.from(this._entitiesSet) as EntityContext[];
     }
diff --git a/src/parser/flink/flinkEntityCollector.ts b/src/parser/flink/flinkEntityCollector.ts
index f39142d3..7bb1c08e 100644
--- a/src/parser/flink/flinkEntityCollector.ts
+++ b/src/parser/flink/flinkEntityCollector.ts
@@ -45,7 +45,17 @@ export class FlinkEntityCollector extends EntityCollector implements FlinkSqlPar
     }
 
     exitTablePath(ctx: TablePathContext) {
-        this.pushEntity(ctx, EntityContextType.TABLE);
+        const needCollectAttr = this.getRootStmt()?.stmtContextType === StmtContextType.SELECT_STMT;
+        this.pushEntity(
+            ctx,
+            EntityContextType.TABLE,
+            needCollectAttr
+                ? {
+                      attrNameList: [AttrName.alias],
+                      endContext: QueryStatementContext.name,
+                  }
+                : undefined
+        );
     }
 
     exitTablePathCreate(ctx: TablePathCreateContext) {
diff --git a/test/parser/flink/contextCollect/entityCollector.test.ts b/test/parser/flink/contextCollect/entityCollector.test.ts
index f57a6c2b..a3d1fd7a 100644
--- a/test/parser/flink/contextCollect/entityCollector.test.ts
+++ b/test/parser/flink/contextCollect/entityCollector.test.ts
@@ -208,6 +208,14 @@ describe('Flink entity collector tests', () => {
 
         expect(tableEntity.entityContextType).toBe(EntityContextType.TABLE);
         expect(tableEntity.text).toBe('Orders');
+        expect(tableEntity[AttrName.alias]).toEqual({
+            text: 'o1',
+            startIndex: 607,
+            endIndex: 608,
+            line: 23,
+            startColumn: 45,
+            endColumn: 47,
+        });
         expect(tableEntity.belongStmt.stmtContextType).toBe(StmtContextType.SELECT_STMT);
         if (isCommonEntityContext(tableEntity)) {
             expect(tableEntity.columns).toBeUndefined();
@@ -360,8 +368,8 @@ describe('Flink entity collector tests', () => {
         expect(allEntities[0].text).toBe('view1');
         expect(allEntities[0][AttrName.comment]).toEqual({
             text: "'this is a view'",
-            startIndex: 1199,
-            endIndex: 1214,
+            startIndex: 1205,
+            endIndex: 1220,
             line: 42,
             startColumn: 39,
             endColumn: 55,
@@ -389,27 +397,27 @@ describe('Flink entity collector tests', () => {
         expect(dbEntity.text).toBe('db1');
         expect(dbEntity[AttrName.comment]).toEqual({
             text: "'this is a created database'",
-            startIndex: 1290,
-            endIndex: 1317,
+            startIndex: 1296,
+            endIndex: 1323,
             line: 44,
             startColumn: 43,
             endColumn: 71,
         });
         expect(dbEntity.position).toEqual({
             endColumn: 34,
-            endIndex: 1280,
+            endIndex: 1286,
             line: 44,
             startColumn: 31,
-            startIndex: 1278,
+            startIndex: 1284,
         });
 
         expect(dbEntity.belongStmt.stmtContextType).toBe(StmtContextType.CREATE_DATABASE_STMT);
         expect(dbEntity.belongStmt.position).toEqual({
             endColumn: 119,
-            endIndex: 1365,
+            endIndex: 1371,
             endLine: 44,
             startColumn: 1,
-            startIndex: 1248,
+            startIndex: 1254,
             startLine: 44,
         });
         if (isCommonEntityContext(dbEntity)) {
@@ -434,10 +442,10 @@ describe('Flink entity collector tests', () => {
         expect(functionEntity.text).toBe('tempFunction');
         expect(functionEntity.position).toEqual({
             endColumn: 43,
-            endIndex: 1410,
+            endIndex: 1416,
             line: 46,
             startColumn: 31,
-            startIndex: 1399,
+            startIndex: 1405,
         });
 
         expect(functionEntity.belongStmt.stmtContextType).toBe(
@@ -445,10 +453,10 @@ describe('Flink entity collector tests', () => {
         );
         expect(functionEntity.belongStmt.position).toEqual({
             endColumn: 58,
-            endIndex: 1425,
+            endIndex: 1431,
             endLine: 46,
             startColumn: 1,
-            startIndex: 1369,
+            startIndex: 1375,
             startLine: 46,
         });
         if (isFuncEntityContext(functionEntity)) {
diff --git a/test/parser/flink/contextCollect/fixtures/common.sql b/test/parser/flink/contextCollect/fixtures/common.sql
index a7dd98cf..ce697f05 100644
--- a/test/parser/flink/contextCollect/fixtures/common.sql
+++ b/test/parser/flink/contextCollect/fixtures/common.sql
@@ -20,7 +20,7 @@ CREATE TABLE Orders_with_watermark (
     INCLUDING GENERATED
 );
 
-SELECT order_id, price + tax FROM Orders;
+SELECT order_id, price + tax FROM Orders AS o1;
 
 SELECT * FROM Orders LEFT JOIN Product ON Orders.product_id = Product.id;