feat: update interface of attrInfo and alter entitycollect ts file

zhaoge · zhaoge · commit 418e38e670fa · 2024-08-02T15:02:08.000+08:00
diff --git a/src/parser/common/entityCollector.ts b/src/parser/common/entityCollector.ts
@@ -128,7 +128,7 @@ export function isColumnEntityContext(entity: EntityContext): entity is ColumnEn
  *  what we need when collect attribute information
  * */
 interface AttrInfo {
-    attrNameList: AttrName[];
+    attrName: AttrName;
     endContextList: string[];
 }
 
@@ -137,7 +137,7 @@ export function toEntityContext(
     type: EntityContextType,
     input: string,
     belongStmt: StmtContext,
-    attrInfo?: AttrInfo
+    attrInfo?: AttrInfo[]
 ): EntityContext | null {
     const word = ctxToText(ctx, input);
     if (!word) return null;
@@ -169,9 +169,9 @@ export function toEntityContext(
             break;
     }
     if (attrInfo) {
-        for (let k = 0; k < attrInfo?.attrNameList?.length; k++) {
-            const attributeName: AttrName = attrInfo?.attrNameList[k];
-            const attrToken = findAttribute(ctx, attributeName, attrInfo?.endContextList);
+        for (let k = 0; k < attrInfo.length; k++) {
+            const attributeName: AttrName = attrInfo[k]?.attrName;
+            const attrToken = findAttribute(ctx, attributeName, attrInfo[k]?.endContextList);
             if (attrToken) {
                 const attrVal: WordRange | TextSlice | null = isToken(attrToken)
                     ? tokenToWord(attrToken, input)
@@ -334,7 +334,7 @@ export abstract class EntityCollector {
         return stmtContext;
     }
 
-    protected pushEntity(ctx: ParserRuleContext, type: EntityContextType, attrInfo?: AttrInfo) {
+    protected pushEntity(ctx: ParserRuleContext, type: EntityContextType, attrInfo?: AttrInfo[]) {
         const entityContext = toEntityContext(
             ctx,
             type,
diff --git a/src/parser/hive/hiveEntityCollector.ts b/src/parser/hive/hiveEntityCollector.ts
@@ -29,10 +29,12 @@ import { EntityContextType } from '../common/types';
 export class HiveEntityCollector extends EntityCollector implements HiveSqlParserListener {
     /** ====== Entity Begin */
     exitTableNameCreate(ctx: TableNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.TABLE_CREATE, {
-            attrNameList: [AttrName.comment],
-            endContextList: [CreateTableStatementContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.TABLE_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [CreateTableStatementContext.name],
+            },
+        ]);
     }
 
     exitTableName(ctx: TableNameContext) {
@@ -41,37 +43,49 @@ export class HiveEntityCollector extends EntityCollector implements HiveSqlParse
             ctx,
             EntityContextType.TABLE,
             needCollectAttr
-                ? {
-                      attrNameList: [AttrName.alias],
-                      endContextList: [TableSourceContext.name, UniqueJoinSourceContext.name],
-                  }
+                ? [
+                      {
+                          attrName: AttrName.alias,
+                          endContextList: [TableSourceContext.name, UniqueJoinSourceContext.name],
+                      },
+                  ]
                 : undefined
         );
     }
 
     exitColumnNameCreate(ctx: ColumnNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE, {
-            attrNameList: [AttrName.comment, AttrName.colType],
-            endContextList: [ColumnNameTypeConstraintContext.name, ColumnNameCommentContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [ColumnNameCommentContext.name],
+            },
+            {
+                attrName: AttrName.colType,
+                endContextList: [ColumnNameTypeConstraintContext.name],
+            },
+        ]);
     }
 
     exitViewNameCreate(ctx: ViewNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.VIEW_CREATE, {
-            attrNameList: [AttrName.comment],
-            endContextList: [CreateViewStatementContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.VIEW_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [CreateViewStatementContext.name],
+            },
+        ]);
     }
 
     exitViewName(ctx: ViewNameContext) {
         this.pushEntity(ctx, EntityContextType.VIEW);
     }
 
     exitDbSchemaNameCreate(ctx: DbSchemaNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.DATABASE_CREATE, {
-            attrNameList: [AttrName.comment],
-            endContextList: [CreateDatabaseStatementContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.DATABASE_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [CreateDatabaseStatementContext.name],
+            },
+        ]);
     }
 
     exitDbSchemaName(ctx: DbSchemaNameContext) {
diff --git a/src/parser/impala/impalaEntityCollector.ts b/src/parser/impala/impalaEntityCollector.ts
@@ -29,10 +29,15 @@ import { EntityContextType } from '../common/types';
 export class ImpalaEntityCollector extends EntityCollector implements ImpalaSqlParserListener {
     /** ===== Entity begin */
     exitTableNameCreate(ctx: TableNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.TABLE_CREATE, {
-            attrNameList: [AttrName.comment],
-            endContextList: [CreateTableSelectContext.name, CreateKuduTableAsSelectContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.TABLE_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [
+                    CreateTableSelectContext.name,
+                    CreateKuduTableAsSelectContext.name,
+                ],
+            },
+        ]);
     }
 
     exitTableNamePath(ctx: TableNamePathContext) {
@@ -41,30 +46,44 @@ export class ImpalaEntityCollector extends EntityCollector implements ImpalaSqlP
             ctx,
             EntityContextType.TABLE,
             needCollectAttr
-                ? {
-                      attrNameList: [AttrName.alias],
-                      endContextList: [SampledRelationContext.name],
-                  }
+                ? [
+                      {
+                          attrName: AttrName.alias,
+                          endContextList: [SampledRelationContext.name],
+                      },
+                  ]
                 : undefined
         );
     }
 
     exitColumnNamePathCreate(ctx: ColumnNamePathCreateContext) {
-        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE, {
-            attrNameList: [AttrName.comment, AttrName.colType],
-            endContextList: [
-                ColumnDefinitionContext.name,
-                KuduTableElementContext.name,
-                ViewColumnItemContext.name,
-            ],
-        });
+        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [
+                    ColumnDefinitionContext.name,
+                    KuduTableElementContext.name,
+                    ViewColumnItemContext.name,
+                ],
+            },
+            {
+                attrName: AttrName.colType,
+                endContextList: [
+                    ColumnDefinitionContext.name,
+                    KuduTableElementContext.name,
+                    ViewColumnItemContext.name,
+                ],
+            },
+        ]);
     }
 
     exitViewNameCreate(ctx: ViewNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.VIEW_CREATE, {
-            attrNameList: [AttrName.comment],
-            endContextList: [CreateViewContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.VIEW_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [CreateViewContext.name],
+            },
+        ]);
     }
 
     exitViewNamePath(ctx: ViewNamePathContext) {
@@ -76,10 +95,12 @@ export class ImpalaEntityCollector extends EntityCollector implements ImpalaSqlP
     }
 
     exitDatabaseNameCreate(ctx: DatabaseNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.DATABASE_CREATE, {
-            attrNameList: [AttrName.comment],
-            endContextList: [CreateSchemaContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.DATABASE_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [CreateSchemaContext.name],
+            },
+        ]);
     }
 
     exitFunctionNameCreate(ctx: FunctionNameCreateContext) {
diff --git a/src/parser/postgresql/postgreEntityCollector.ts b/src/parser/postgresql/postgreEntityCollector.ts
@@ -44,10 +44,12 @@ export class PostgreSqlEntityCollector extends EntityCollector implements Postgr
             ctx,
             EntityContextType.TABLE,
             needCollectAttr
-                ? {
-                      attrNameList: [AttrName.alias],
-                      endContextList: [Table_refContext.name],
-                  }
+                ? [
+                      {
+                          attrName: AttrName.alias,
+                          endContextList: [Table_refContext.name],
+                      },
+                  ]
                 : undefined
         );
     }
@@ -62,10 +64,12 @@ export class PostgreSqlEntityCollector extends EntityCollector implements Postgr
             ctx,
             EntityContextType.VIEW,
             needCollectAttr
-                ? {
-                      attrNameList: [AttrName.alias],
-                      endContextList: [Table_refContext.name],
-                  }
+                ? [
+                      {
+                          attrName: AttrName.alias,
+                          endContextList: [Table_refContext.name],
+                      },
+                  ]
                 : undefined
         );
     }
@@ -79,10 +83,16 @@ export class PostgreSqlEntityCollector extends EntityCollector implements Postgr
     }
 
     exitColumnNameCreate(ctx: ColumnNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE, {
-            attrNameList: [AttrName.comment, AttrName.colType],
-            endContextList: [Column_defContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [Column_defContext.name],
+            },
+            {
+                attrName: AttrName.colType,
+                endContextList: [Column_defContext.name],
+            },
+        ]);
     }
 
     exitProcedureName(ctx: ProcedureNameContext) {
diff --git a/src/parser/spark/sparkEntityCollector.ts b/src/parser/spark/sparkEntityCollector.ts
@@ -33,10 +33,12 @@ export class SparkEntityCollector extends EntityCollector implements SparkSqlPar
     }
 
     exitNamespaceNameCreate(ctx: NamespaceNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.DATABASE_CREATE, {
-            attrNameList: [AttrName.comment],
-            endContextList: [CreateNamespaceContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.DATABASE_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [CreateNamespaceContext.name],
+            },
+        ]);
     }
 
     exitTableName(ctx: TableNameContext) {
@@ -45,44 +47,56 @@ export class SparkEntityCollector extends EntityCollector implements SparkSqlPar
             ctx,
             EntityContextType.TABLE,
             needCollectAttr
-                ? {
-                      attrNameList: [AttrName.alias],
-                      endContextList: [RelationPrimaryContext.name],
-                  }
+                ? [
+                      {
+                          attrName: AttrName.alias,
+                          endContextList: [RelationPrimaryContext.name],
+                      },
+                  ]
                 : undefined
         );
     }
 
     exitTableNameCreate(ctx: TableNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.TABLE_CREATE, {
-            attrNameList: [AttrName.comment],
-            endContextList: [CreateTableContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.TABLE_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [CreateTableContext.name],
+            },
+        ]);
     }
 
     exitViewName(ctx: ViewNameContext) {
         this.pushEntity(ctx, EntityContextType.VIEW);
     }
 
     exitViewNameCreate(ctx: ViewNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.VIEW_CREATE, {
-            attrNameList: [AttrName.comment],
-            endContextList: [CreateViewContext.name],
-        });
+        this.pushEntity(ctx, EntityContextType.VIEW_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [CreateViewContext.name],
+            },
+        ]);
     }
 
     exitFunctionNameCreate(ctx: FunctionNameCreateContext) {
         this.pushEntity(ctx, EntityContextType.FUNCTION_CREATE);
     }
 
     exitColumnNameCreate(ctx: ColumnNameCreateContext) {
-        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE, {
-            attrNameList: [AttrName.comment, AttrName.colType],
-            endContextList: [
-                CreateOrReplaceTableColTypeContext.name,
-                IdentifierCommentContext.name,
-            ],
-        });
+        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE, [
+            {
+                attrName: AttrName.comment,
+                endContextList: [
+                    CreateOrReplaceTableColTypeContext.name,
+                    IdentifierCommentContext.name,
+                ],
+            },
+            {
+                attrName: AttrName.colType,
+                endContextList: [CreateOrReplaceTableColTypeContext.name],
+            },
+        ]);
     }
 
     /** ===== Statement begin */
diff --git a/src/parser/trino/trinoEntityCollector.ts b/src/parser/trino/trinoEntityCollector.ts
