feat(pgsql): pgsql collect entity (#268)

* feat(pgsql): pgsql collect entity * feat(pgsql): optimize some name --------- Co-authored-by: zhaoge <>
DTStack · Feb 29, 2024 · 2ae10b8 · 2ae10b8
1 parent d149072
commit 2ae10b8
Show file tree

Hide file tree

Showing 12 changed files with 17,186 additions and 16,101 deletions.
diff --git a/src/grammar/pgsql/PostgreSQLParser.g4 b/src/grammar/pgsql/PostgreSQLParser.g4
diff --git a/src/lib/pgsql/PostgreSQLParser.interp b/src/lib/pgsql/PostgreSQLParser.interp
diff --git a/src/lib/pgsql/PostgreSQLParser.ts b/src/lib/pgsql/PostgreSQLParser.ts
diff --git a/src/lib/pgsql/PostgreSQLParserListener.ts b/src/lib/pgsql/PostgreSQLParserListener.ts
diff --git a/src/lib/pgsql/PostgreSQLParserVisitor.ts b/src/lib/pgsql/PostgreSQLParserVisitor.ts
diff --git a/src/parser/pgsql.ts → src/parser/pgsql/index.ts b/src/parser/pgsql.ts → src/parser/pgsql/index.ts
@@ -1,11 +1,15 @@
-import { Token } from 'antlr4ng';
 import { CandidatesCollection } from 'antlr4-c3';
-import { PostgreSQLLexer } from '../lib/pgsql/PostgreSQLLexer';
-import { PostgreSQLParser, ProgramContext, SingleStmtContext } from '../lib/pgsql/PostgreSQLParser';
-import BasicParser from './common/basicParser';
-import { PostgreSQLParserListener } from '../lib/pgsql/PostgreSQLParserListener';
-import { EntityContextType, Suggestions, SyntaxSuggestion } from './common/basic-parser-types';
-import { StmtContextType } from './common/entityCollector';
+import { Token } from 'antlr4ng';
+
+import { PostgreSQLLexer } from '../../lib/pgsql/PostgreSQLLexer';
+import { PostgreSQLParser, ProgramContext } from '../../lib/pgsql/PostgreSQLParser';
+import { EntityContextType, Suggestions, SyntaxSuggestion } from '../common/basic-parser-types';
+import BasicParser from '../common/basicParser';
+import { StmtContextType } from '../common/entityCollector';
+import PostgreSQLEntityCollector from './postgreEntityCollector';
+import PostgreSqlSplitListener from './postgreSplitListener';
+
+export { PostgreSQLEntityCollector, PostgreSqlSplitListener };
 
 export default class PostgresSQL extends BasicParser<
     PostgreSQLLexer,
@@ -39,7 +43,7 @@ export default class PostgresSQL extends BasicParser<
     ]);
 
     protected get splitListener() {
-        return new PgSqlSplitListener();
+        return new PostgreSqlSplitListener();
     }
 
     protected processCandidates(
@@ -145,20 +149,3 @@ export default class PostgresSQL extends BasicParser<
         };
     }
 }
-
-export class PgSqlSplitListener implements PostgreSQLParserListener {
-    private _statementsContext: SingleStmtContext[] = [];
-
-    exitSingleStmt = (ctx: SingleStmtContext) => {
-        this._statementsContext.push(ctx);
-    };
-
-    visitTerminal() {}
-    visitErrorNode() {}
-    enterEveryRule() {}
-    exitEveryRule() {}
-
-    get statementsContext() {
-        return this._statementsContext;
-    }
-}
diff --git a/src/parser/pgsql/postgreEntityCollector.ts b/src/parser/pgsql/postgreEntityCollector.ts
@@ -0,0 +1,209 @@
+import type {
+    ColumnCreateTableContext,
+    ColumnNameCreateContext,
+    CreateDatabaseContext,
+    CreateForeignTableContext,
+    CreateMaterializedViewContext,
+    CreatePartitionForeignTableContext,
+    CreateViewContext,
+    DatabaseNameContext,
+    DatabaseNameCreateContext,
+    FunctionNameCreateContext,
+    InsertStatementContext,
+    QueryCreateTableContext,
+    SelectStatementContext,
+    SingleStmtContext,
+    TableNameContext,
+    TableNameCreateContext,
+    ViewNameContext,
+    ViewNameCreateContext,
+} from '../../lib/pgsql/PostgreSQLParser';
+import type { PostgreSQLParserListener } from '../../lib/pgsql/PostgreSQLParserListener';
+import { EntityContextType } from '../common/basic-parser-types';
+import EntityCollector, {
+    EntityContext,
+    StmtContext,
+    StmtContextType,
+} from '../common/entityCollector';
+
+export default class PostgreSQLEntityCollector
+    extends EntityCollector
+    implements PostgreSQLParserListener
+{
+    visitTerminal() {}
+    visitErrorNode() {}
+    enterEveryRule() {}
+    exitEveryRule() {}
+    combineRootStmtEntities(
+        stmtContext: StmtContext,
+        entitiesInsideStmt: EntityContext[]
+    ): EntityContext[] {
+        if (
+            stmtContext.stmtContextType === StmtContextType.CREATE_TABLE_STMT ||
+            stmtContext.stmtContextType === StmtContextType.CREATE_VIEW_STMT
+        ) {
+            return this.combineCreateTableOrViewStmtEntities(stmtContext, entitiesInsideStmt);
+        }
+        return entitiesInsideStmt;
+    }
+
+    combineCreateTableOrViewStmtEntities(
+        stmtContext: StmtContext,
+        entitiesInsideStmt: EntityContext[]
+    ): EntityContext[] {
+        const columns: EntityContext[] = [];
+        const relatedEntities: EntityContext[] = [];
+        let mainEntity: EntityContext = null;
+        const finalEntities = entitiesInsideStmt.reduce((result, entity) => {
+            if (entity.belongStmt !== stmtContext) {
+                if (
+                    entity.entityContextType !== EntityContextType.COLUMN &&
+                    entity.entityContextType !== EntityContextType.COLUMN_CREATE
+                ) {
+                    relatedEntities.push(entity);
+                    result.push(entity);
+                }
+                return result;
+            }
+
+            if (entity.entityContextType === EntityContextType.COLUMN_CREATE) {
+                columns.push(entity);
+            } else if (
+                entity.entityContextType === EntityContextType.TABLE_CREATE ||
+                entity.entityContextType === EntityContextType.VIEW_CREATE
+            ) {
+                mainEntity = entity;
+                result.push(entity);
+                return result;
+            } else if (entity.entityContextType !== EntityContextType.COLUMN) {
+                relatedEntities.push(entity);
+                result.push(entity);
+            }
+            return result;
+        }, []);
+
+        if (columns.length) {
+            mainEntity.columns = columns;
+        }
+
+        if (relatedEntities.length) {
+            mainEntity.relatedEntities = relatedEntities;
+        }
+
+        return finalEntities;
+    }
+
+    /** ====== Entity Begin */
+    exitDatabaseName(ctx: DatabaseNameContext) {
+        this.pushEntity(ctx, EntityContextType.DATABASE);
+    }
+
+    exitDatabaseNameCreate(ctx: DatabaseNameCreateContext) {
+        this.pushEntity(ctx, EntityContextType.DATABASE_CREATE);
+    }
+
+    exitTableName(ctx: TableNameContext) {
+        this.pushEntity(ctx, EntityContextType.TABLE);
+    }
+
+    exitTableNameCreate(ctx: TableNameCreateContext) {
+        this.pushEntity(ctx, EntityContextType.TABLE_CREATE);
+    }
+
+    exitViewName(ctx: ViewNameContext) {
+        this.pushEntity(ctx, EntityContextType.VIEW);
+    }
+
+    exitViewNameCreate(ctx: ViewNameCreateContext) {
+        this.pushEntity(ctx, EntityContextType.VIEW_CREATE);
+    }
+
+    exitFunctionNameCreate(ctx: FunctionNameCreateContext) {
+        this.pushEntity(ctx, EntityContextType.FUNCTION_CREATE);
+    }
+
+    exitColumnNameCreate(ctx: ColumnNameCreateContext) {
+        this.pushEntity(ctx, EntityContextType.COLUMN_CREATE);
+    }
+
+    /** ===== Statement begin */
+    enterSingleStatement(ctx: SingleStmtContext) {
+        this.pushStmt(ctx, StmtContextType.COMMON_STMT);
+    }
+
+    exitSingleStatement(ctx: SingleStmtContext) {
+        this.popStmt();
+    }
+
+    enterCreateDatabase(ctx: CreateDatabaseContext) {
+        this.pushStmt(ctx, StmtContextType.CREATE_DATABASE_STMT);
+    }
+
+    exitCreateDatabase(ctx: CreateDatabaseContext) {
+        this.popStmt();
+    }
+
+    enterQueryCreateTable(ctx: QueryCreateTableContext) {
+        this.pushStmt(ctx, StmtContextType.CREATE_TABLE_STMT);
+    }
+
+    exitQueryCreateTable(ctx: QueryCreateTableContext) {
+        this.popStmt();
+    }
+
+    enterColumnCreateTable(ctx: ColumnCreateTableContext) {
+        this.pushStmt(ctx, StmtContextType.CREATE_TABLE_STMT);
+    }
+
+    exitColumnCreateTable(ctx: ColumnCreateTableContext) {
+        this.popStmt();
+    }
+
+    enterCreateForeignTable(ctx: CreateForeignTableContext) {
+        this.pushStmt(ctx, StmtContextType.CREATE_TABLE_STMT);
+    }
+
+    exitCreateForeignTable(ctx: CreateForeignTableContext) {
+        this.popStmt();
+    }
+
+    enterCreatePartitionForeignTable(ctx: CreatePartitionForeignTableContext) {
+        this.pushStmt(ctx, StmtContextType.CREATE_TABLE_STMT);
+    }
+
+    exitCreatePartitionForeignTable(ctx: CreatePartitionForeignTableContext) {
+        this.popStmt();
+    }
+
+    enterCreateView(ctx: CreateViewContext) {
+        this.pushStmt(ctx, StmtContextType.CREATE_VIEW_STMT);
+    }
+
+    exitCreateView(ctx: CreateViewContext) {
+        this.popStmt();
+    }
+
+    enterCreateMaterializedView(ctx: CreateMaterializedViewContext) {
+        this.pushStmt(ctx, StmtContextType.CREATE_VIEW_STMT);
+    }
+
+    exitCreateMaterializedView(ctx: CreateMaterializedViewContext) {
+        this.popStmt();
+    }
+
+    enterSelectStatement(ctx: SelectStatementContext) {
+        this.pushStmt(ctx, StmtContextType.SELECT_STMT);
+    }
+
+    exitSelectStatement(ctx: SelectStatementContext) {
+        this.popStmt();
+    }
+
+    enterInsertStatement(ctx: InsertStatementContext) {
+        this.pushStmt(ctx, StmtContextType.INSERT_STMT);
+    }
+
+    exitInsertStatement(ctx: InsertStatementContext) {
+        this.popStmt();
+    }
+}
diff --git a/src/parser/pgsql/postgreSplitListener.ts b/src/parser/pgsql/postgreSplitListener.ts
@@ -0,0 +1,20 @@
+import { SingleStmtContext } from '../../lib/pgsql/PostgreSQLParser';
+import { PostgreSQLParserListener } from '../../lib/pgsql/PostgreSQLParserListener';
+
+export default class PostgreSqlSplitListener implements PostgreSQLParserListener {
+    visitTerminal() {}
+    visitErrorNode() {}
+    enterEveryRule() {}
+    exitEveryRule() {}
+    private _statementsContext: SingleStmtContext[] = [];
+
+    exitSingleStmt = (ctx: SingleStmtContext) => {
+        this._statementsContext.push(ctx);
+    };
+
+    enterSingleStmt = (ctx: SingleStmtContext) => {};
+
+    get statementsContext() {
+        return this._statementsContext;
+    }
+}