feat(hivesql): collect hive sql's attribute(comment,alias,colType)

DTStack · Jul 12, 2024 · 2deea8f · 2deea8f
1 parent 8b55899
commit 2deea8f
Show file tree

Hide file tree

Showing 5 changed files with 297 additions and 20 deletions.
diff --git a/src/grammar/hive/HiveSqlParser.g4 b/src/grammar/hive/HiveSqlParser.g4
@@ -887,7 +887,9 @@ tableConstraint
     ;
 
 columnNameTypeConstraint
-    : colName=columnNameCreate columnType columnConstraint? (KW_COMMENT comment=StringLiteral)?
+    : colName=columnNameCreate colType=columnType columnConstraint? (
+        KW_COMMENT comment=StringLiteral
+    )?
     ;
 
 columnConstraint
@@ -1363,10 +1365,12 @@ joinToken
     ;
 
 lateralView
-    : KW_LATERAL KW_VIEW KW_OUTER function_ tableAlias (KW_AS id_ (COMMA id_)*)?
+    : KW_LATERAL KW_VIEW KW_OUTER function_ alias=tableAlias (KW_AS id_ (COMMA id_)*)?
     | COMMA? KW_LATERAL (
-        KW_VIEW function_ tableAlias (KW_AS id_ (COMMA id_)*)?
-        | KW_TABLE LPAREN valuesClause RPAREN KW_AS? tableAlias (LPAREN id_ (COMMA id_)* RPAREN)?
+        KW_VIEW function_ alias=tableAlias (KW_AS id_ (COMMA id_)*)?
+        | KW_TABLE LPAREN valuesClause RPAREN KW_AS? alias=tableAlias (
+            LPAREN id_ (COMMA id_)* RPAREN
+        )?
     )
     ;
 

diff --git a/src/lib/hive/HiveSqlParser.ts b/src/lib/hive/HiveSqlParser.ts
@@ -10836,7 +10836,7 @@ export class HiveSqlParser extends SQLParserBase {
             this.state = 2290;
             localContext._colName = this.columnNameCreate();
             this.state = 2291;
-            this.columnType();
+            localContext._colType = this.columnType();
             this.state = 2293;
             this.errorHandler.sync(this);
             _la = this.tokenStream.LA(1);
@@ -16816,7 +16816,7 @@ export class HiveSqlParser extends SQLParserBase {
                 this.state = 3480;
                 this.function_();
                 this.state = 3481;
-                this.tableAlias();
+                localContext._alias = this.tableAlias();
                 this.state = 3491;
                 this.errorHandler.sync(this);
                 _la = this.tokenStream.LA(1);
@@ -16874,7 +16874,7 @@ export class HiveSqlParser extends SQLParserBase {
                     this.state = 3498;
                     this.function_();
                     this.state = 3499;
-                    this.tableAlias();
+                    localContext._alias = this.tableAlias();
                     this.state = 3509;
                     this.errorHandler.sync(this);
                     _la = this.tokenStream.LA(1);
@@ -16928,7 +16928,7 @@ export class HiveSqlParser extends SQLParserBase {
                     }
 
                     this.state = 3518;
-                    this.tableAlias();
+                    localContext._alias = this.tableAlias();
                     this.state = 3530;
                     this.errorHandler.sync(this);
                     switch (this.interpreter.adaptivePredict(this.tokenStream, 458, this.context) ) {
@@ -34747,16 +34747,17 @@ export class TableConstraintContext extends antlr.ParserRuleContext {
 
 export class ColumnNameTypeConstraintContext extends antlr.ParserRuleContext {
     public _colName?: ColumnNameCreateContext;
+    public _colType?: ColumnTypeContext;
     public _comment?: Token | null;
     public constructor(parent: antlr.ParserRuleContext | null, invokingState: number) {
         super(parent, invokingState);
     }
-    public columnType(): ColumnTypeContext {
-        return this.getRuleContext(0, ColumnTypeContext)!;
-    }
     public columnNameCreate(): ColumnNameCreateContext {
         return this.getRuleContext(0, ColumnNameCreateContext)!;
     }
+    public columnType(): ColumnTypeContext {
+        return this.getRuleContext(0, ColumnTypeContext)!;
+    }
     public columnConstraint(): ColumnConstraintContext | null {
         return this.getRuleContext(0, ColumnConstraintContext);
     }
@@ -38744,6 +38745,7 @@ export class JoinTokenContext extends antlr.ParserRuleContext {
 
 
 export class LateralViewContext extends antlr.ParserRuleContext {
+    public _alias?: TableAliasContext;
     public constructor(parent: antlr.ParserRuleContext | null, invokingState: number) {
         super(parent, invokingState);
     }

diff --git a/src/parser/hive/hiveEntityCollector.ts b/src/parser/hive/hiveEntityCollector.ts
@@ -1,12 +1,13 @@
-import { EntityContextType } from '../common/types';
 import { HiveSqlParserListener } from '../../lib';
 import {
     ColumnNameCreateContext,
+    ColumnNameTypeConstraintContext,
     CreateDatabaseStatementContext,
     CreateFunctionStatementContext,
     CreateMaterializedViewStatementContext,
     CreateTableStatementContext,
     CreateViewStatementContext,
+    DatabaseCommentContext,
     DbSchemaNameContext,
     DbSchemaNameCreateContext,
     FromInsertStmtContext,
@@ -15,37 +16,63 @@ import {
     InsertStmtContext,
     SelectStatementContext,
     StatementContext,
+    TableCommentContext,
     TableNameContext,
     TableNameCreateContext,
+    TableSourceContext,
+    UniqueJoinSourceContext,
     ViewNameContext,
     ViewNameCreateContext,
 } from '../../lib/hive/HiveSqlParser';
-import { StmtContextType, EntityCollector } from '../common/entityCollector';
+import { AttrName, EntityCollector, StmtContextType } from '../common/entityCollector';
+import { EntityContextType } from '../common/types';
 
 export class HiveEntityCollector extends EntityCollector implements HiveSqlParserListener {
     /** ====== Entity Begin */
     exitTableNameCreate(ctx: TableNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.TABLE_CREATE);
+        this.pushEntity(ctx, EntityContextType.TABLE_CREATE, {
+            attrNameList: [AttrName.comment],
+            endContextList: [TableCommentContext.name],
+        });
     }
 
     exitTableName(ctx: TableNameContext) {
-        this.pushEntity(ctx, EntityContextType.TABLE);
+        const needCollectAttr = this.getRootStmt()?.stmtContextType === StmtContextType.SELECT_STMT;
+        this.pushEntity(
+            ctx,
+            EntityContextType.TABLE,
+            needCollectAttr
+                ? {
+                      attrNameList: [AttrName.alias],
+                      endContextList: [TableSourceContext.name, UniqueJoinSourceContext.name],
+                  }
+                : undefined
+        );
     }
 
     exitColumnNameCreate(ctx: ColumnNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE);
+        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE, {
+            attrNameList: [AttrName.comment, AttrName.colType],
+            endContextList: [ColumnNameTypeConstraintContext.name],
+        });
     }
 
     exitViewNameCreate(ctx: ViewNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.VIEW_CREATE);
+        this.pushEntity(ctx, EntityContextType.VIEW_CREATE, {
+            attrNameList: [AttrName.comment],
+            endContextList: [TableCommentContext.name],
+        });
     }
 
     exitViewName(ctx: ViewNameContext) {
         this.pushEntity(ctx, EntityContextType.VIEW);
     }
 
     exitDbSchemaNameCreate(ctx: DbSchemaNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.DATABASE_CREATE);
+        this.pushEntity(ctx, EntityContextType.DATABASE_CREATE, {
+            attrNameList: [AttrName.comment],
+            endContextList: [DatabaseCommentContext.name],
+        });
     }
 
     exitDbSchemaName(ctx: DbSchemaNameContext) {