From e7729dae33b6eab7c75eb0b48a9d160e22a17b88 Mon Sep 17 00:00:00 2001
From: Kijin-Seija <almon123@qq.com>
Date: Wed, 10 Apr 2024 18:40:49 +0800
Subject: [PATCH] =?UTF-8?q?0.0.1-alpha.5=20-=20entity=E6=B7=BB=E5=8A=A0bel?=
 =?UTF-8?q?ongsToEntity=E5=B1=9E=E6=80=A7=20-=20=E4=BF=AE=E6=94=B9entity?=
 =?UTF-8?q?=E5=92=8Cstmt=E7=9A=84=E5=85=B3=E8=81=94=E9=80=BB=E8=BE=91,?=
 =?UTF-8?q?=E4=B8=8D=E5=86=8D=E5=8F=96=E6=A0=88=E9=A1=B6=E5=85=83=E7=B4=A0?=
 =?UTF-8?q?,=E8=80=8C=E6=98=AF=E6=90=9C=E7=B4=A2=E5=8C=B9=E9=85=8DruleInde?=
 =?UTF-8?q?x=E7=9A=84=E5=85=83=E7=B4=A0.=E8=BF=99=E6=A0=B7=E5=8F=AF?=
 =?UTF-8?q?=E4=BB=A5=E5=9C=A8=E5=8C=B9=E9=85=8D=E5=A4=9A=E4=B8=AA=E8=A7=84?=
 =?UTF-8?q?=E5=88=99=E9=93=BE=E6=97=B6=E9=83=BD=E5=85=B3=E8=81=94?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 package.json               |  2 +-
 src/parse/default-rules.ts | 11 +++++++++++
 src/parse/visitor.ts       | 23 ++++++++++++++---------
 src/types.ts               |  3 ++-
 4 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/package.json b/package.json
index a131854..d5e78f3 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "dt-sql-parser-semantic-analyse-plugin",
-  "version": "0.0.1-alpha.4",
+  "version": "0.0.1-alpha.5",
   "description": "an dt-sql-parser plugin with semantic result",
   "type": "module",
   "files": [
diff --git a/src/parse/default-rules.ts b/src/parse/default-rules.ts
index f130de3..9771302 100644
--- a/src/parse/default-rules.ts
+++ b/src/parse/default-rules.ts
@@ -1,6 +1,7 @@
 import { PostgreSQLParser } from 'dt-sql-parser/dist/lib/pgsql/PostgreSQLParser'
 
 export const defaultStmts = [
+  'stmt',
   // select statement
   'selectstmt'
 ]
@@ -21,6 +22,16 @@ export const defaultEntities = [
 ]
 
 export const defaultRules: Record<string, number[]> = {
+  // 通用的简单column规则（不带.运算符的column)
+  common_column_simple: [
+    PostgreSQLParser.RULE_stmt,
+    PostgreSQLParser.RULE_column_name
+  ],
+  // 通用的复合column规则（带.运算符的column)
+  common_column_ref: [
+    PostgreSQLParser.RULE_stmt,
+    PostgreSQLParser.RULE_columnref
+  ],
   select_target_column_simple: [
     PostgreSQLParser.RULE_selectstmt,
     PostgreSQLParser.RULE_select_clause,
diff --git a/src/parse/visitor.ts b/src/parse/visitor.ts
index dcedbec..e5f08e7 100644
--- a/src/parse/visitor.ts
+++ b/src/parse/visitor.ts
@@ -58,26 +58,31 @@ export class SQLVisitor extends AbstractParseTreeVisitor<void> implements Postgr
     (this as any)[visitorName] = (ctx: ParserRuleContext) => {
       const chain = this.getNodeChain(ctx)
       for (const rule of rules) {
-        if (this.matchRules(chain, this.rules.get(rule))) {
+        const ruleChain = this.rules.get(rule)
+        if (!ruleChain) continue
+        if (this.matchRules(chain, ruleChain)) {
+          const ruleChainBegin = ruleChain[0]
+          const beginStmt = this.stmtStack.find(stmt => stmt.type === ruleChainBegin)
+          const beginEntity = this.entityStack.find(entity => entity.type === ruleChainBegin)
           const result: Entity = {
             rule,
             text: ctx.text,
             type: ctx.ruleIndex,
             caret: withCaret(ctx),
-            belongsToStmt: this.stmtStack[this.stmtStack.length - 1],
+            belongsToStmt: beginStmt || null,
+            belongsToEntity: beginEntity || null,
             relatedEntities: {}
           }
-          if (this.entityStack[this.entityStack.length - 1]) {
-            if (!this.entityStack[this.entityStack.length - 1].relatedEntities[rule]) this.entityStack[this.entityStack.length - 1].relatedEntities[rule] = []
-            this.entityStack[this.entityStack.length - 1].relatedEntities[rule].push(result)
-          } else {
-            if (!this.stmtStack[this.stmtStack.length - 1].relatedEntities[rule]) this.stmtStack[this.stmtStack.length - 1].relatedEntities[rule] = []
-            this.stmtStack[this.stmtStack.length - 1].relatedEntities[rule].push(result)
+          if (beginEntity) {
+            if (!beginEntity.relatedEntities[rule]) beginEntity.relatedEntities[rule] = []
+            beginEntity.relatedEntities[rule].push(result)
+          } else if (beginStmt) {
+            if (!beginStmt.relatedEntities[rule]) beginStmt.relatedEntities[rule] = []
+            beginStmt.relatedEntities[rule].push(result)
           }
           if (withCaret(ctx)) this.result.nerestCaretEntityList.push(result)
           this.entityStack.push(result)
           isHitRule = true
-          break
         }
       }
       this.visitChildren(ctx)
diff --git a/src/types.ts b/src/types.ts
index 66f67f4..5d10762 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -52,7 +52,8 @@ export interface Entity {
   text: string
   type: number
   caret: boolean
-  belongsToStmt: Stmt
+  belongsToStmt: Stmt | null
+  belongsToEntity: Entity | null
   relatedEntities: Record<string, Entity[]>
 }