add TableColumnAliasContext

Author: xlorne

springboot-starter-data-authorization/src/main/java/com/codingapi/springboot/authorization/enhancer/DataPermissionSQLEnhancer.java

@@ -4,18 +4,18 @@
 import com.codingapi.springboot.authorization.handler.Condition;
 import com.codingapi.springboot.authorization.handler.RowHandler;
 import lombok.Getter;
+import net.sf.jsqlparser.expression.Alias;
 import net.sf.jsqlparser.expression.Expression;
 import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
 import net.sf.jsqlparser.parser.CCJSqlParserUtil;
+import net.sf.jsqlparser.schema.Column;
 import net.sf.jsqlparser.schema.Table;
 import net.sf.jsqlparser.statement.Statement;
-import net.sf.jsqlparser.statement.select.FromItem;
-import net.sf.jsqlparser.statement.select.Join;
-import net.sf.jsqlparser.statement.select.PlainSelect;
-import net.sf.jsqlparser.statement.select.Select;
+import net.sf.jsqlparser.statement.select.*;
 
 import java.sql.SQLException;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -29,12 +29,15 @@ public class DataPermissionSQLEnhancer {
     @Getter
     private final Map<String, String> tableAlias;
 
+    private final TableColumnAliasContext aliasContext;
+
     // 构造函数
     public DataPermissionSQLEnhancer(String sql, RowHandler rowHandler) {
         // 如何sql中存在? 则在?后面添加空格
         this.sql = sql.replaceAll("\\?", " ? ");
         this.rowHandler = rowHandler;
         this.tableAlias = new HashMap<>();
+        this.aliasContext = new TableColumnAliasContext();
     }
 
     // 获取增强后的SQL
@@ -44,8 +47,9 @@ public String getNewSQL() throws SQLException {
             if (statement instanceof Select) {
                 Select select = (Select) statement;
                 PlainSelect plainSelect = select.getPlainSelect();
-
                 this.enhanceDataPermissionInSelect(plainSelect);
+                this.appendColumnAlias(plainSelect.getSelectItems());
+                aliasContext.print();
                 System.out.println(tableAlias);
                 return statement.toString();
             }
@@ -63,11 +67,13 @@ private void enhanceDataPermissionInSelect(PlainSelect plainSelect) throws Excep
         // FROM 项是表
         if (fromItem instanceof Table) {
             Table table = (Table) fromItem;
+            this.appendTableAlias(fromItem);
             this.injectDataPermissionCondition(plainSelect, table, plainSelect.getWhere());
         }
 
         // FROM是子查询
         if (fromItem instanceof Select) {
+            this.appendTableAlias(fromItem);
             PlainSelect subPlainSelect = ((Select) fromItem).getPlainSelect();
             this.enhanceDataPermissionInSelect(subPlainSelect);
         }
@@ -77,15 +83,60 @@ private void enhanceDataPermissionInSelect(PlainSelect plainSelect) throws Excep
             for (Join join : plainSelect.getJoins()) {
                 if (join.getRightItem() instanceof Select) {
                     PlainSelect subPlainSelect = ((Select) join.getRightItem()).getPlainSelect();
+                    this.appendTableAlias(join.getRightItem());
                     this.enhanceDataPermissionInSelect(subPlainSelect);
                 }
                 if (join.getRightItem() instanceof Table) {
+                    this.appendTableAlias(join.getRightItem());
                     injectDataPermissionCondition(plainSelect, (Table) join.getRightItem(), plainSelect.getWhere());
                 }
             }
         }
     }
 
+
+    private void appendTableAlias(FromItem fromItem) {
+        if (fromItem instanceof Table) {
+            Table table = (Table) fromItem;
+            Alias alias = table.getAlias();
+            String aliasName = alias.getName();
+            aliasContext.add(aliasName, table.getName());
+        }
+        if (fromItem instanceof Select) {
+            Select select = (Select) fromItem;
+            PlainSelect plainSelect = select.getPlainSelect();
+            this.appendTableAlias(plainSelect.getFromItem());
+            List<Join> joins = plainSelect.getJoins();
+            if (joins != null) {
+                for (Join join : joins) {
+                    if (join.getRightItem() instanceof Table) {
+                        this.appendTableAlias(join.getRightItem());
+                    }
+                    if (join.getRightItem() instanceof Select) {
+                        this.appendTableAlias(join.getRightItem());
+                    }
+                }
+            }
+            this.appendColumnAlias(plainSelect.getSelectItems());
+        }
+
+    }
+
+
+    private void appendColumnAlias(List<SelectItem<?>> selectItems) {
+        if (selectItems != null) {
+            for (SelectItem<?> selectItem : selectItems) {
+                if (selectItem.getExpression() instanceof Column) {
+                    Column column = (Column) selectItem.getExpression();
+                    String aliasName = column.getTable().getName();
+                    String columnName = column.getColumnName();
+                    aliasContext.add(aliasName, columnName);
+                }
+            }
+        }
+    }
+
+
     // 注入数据权限条件
     private void injectDataPermissionCondition(PlainSelect plainSelect, Table table, Expression where) throws Exception {
         String tableName = table.getName();

springboot-starter-data-authorization/src/main/java/com/codingapi/springboot/authorization/enhancer/TableColumnAlias.java

@@ -0,0 +1,28 @@
+package com.codingapi.springboot.authorization.enhancer;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class TableColumnAlias {
+
+    private final String aliasName;
+    private final String columnName;
+
+    private Map<String,TableColumnAlias> children;
+
+    public String getKey(){
+        return aliasName+"."+columnName;
+    }
+
+    public TableColumnAlias(String aliasName, String columnName) {
+        this.aliasName = aliasName;
+        this.columnName = columnName;
+        this.children = new HashMap<>();
+    }
+
+    public TableColumnAlias add(String aliasName,String columnName){
+        TableColumnAlias tableColumnAlias = new TableColumnAlias(aliasName,columnName);
+        children.put(tableColumnAlias.getKey(),tableColumnAlias);
+        return tableColumnAlias;
+    }
+}

springboot-starter-data-authorization/src/main/java/com/codingapi/springboot/authorization/enhancer/TableColumnAliasContext.java

@@ -0,0 +1,27 @@
+package com.codingapi.springboot.authorization.enhancer;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class TableColumnAliasContext {
+
+    private final Map<String,TableColumnAlias> tableColumnAliasMap;
+
+    public TableColumnAliasContext() {
+        this.tableColumnAliasMap = new HashMap<>();
+    }
+
+    public TableColumnAlias add(String aliasName,String columnName){
+        TableColumnAlias tableColumnAlias = new TableColumnAlias(aliasName,columnName);
+        tableColumnAliasMap.put(tableColumnAlias.getKey(),tableColumnAlias);
+        return tableColumnAlias;
+    }
+
+
+    public void print(){
+        for(String key:tableColumnAliasMap.keySet()){
+            TableColumnAlias tableColumnAlias = tableColumnAliasMap.get(key);
+            System.out.println(tableColumnAlias.getKey());
+        }
+    }
+}

springboot-starter-data-authorization/src/test/java/com/codingapi/springboot/authorization/DataAuthorizationContextTest.java

@@ -279,8 +279,8 @@ public boolean supportColumnAuthorization(String tableName, String columnName, O
     }
 
 
-//    @Test
-//    @Order(4)
+    @Test
+    @Order(4)
     void test4() throws Exception{
         String sql = "SELECT\n" +
                 "\tUNYiV.id AS '历史工作经历编号',\n" +
@@ -345,7 +345,7 @@ public Condition rowAuthorization(String tableName, String tableAlias) {
 
             @Override
             public <T> T columnAuthorization(String tableName, String columnName, T value) {
-                System.out.println("tableName:" + tableName + ",columnName:" + columnName + ",value:" + value);
+//                System.out.println("tableName:" + tableName + ",columnName:" + columnName + ",value:" + value);
                 return value;
             }
 

springboot-starter-data-authorization/src/test/resources/application.properties

@@ -1,13 +1,13 @@
 
-spring.datasource.driver-class-name=com.codingapi.springboot.authorization.jdbc.AuthorizationJdbcDriver
-spring.datasource.url=jdbc:h2:file:./test.db
-spring.jpa.database-platform=org.hibernate.dialect.H2Dialect
-spring.jpa.hibernate.ddl-auto=create-drop
-spring.jpa.show-sql=true
-
 #spring.datasource.driver-class-name=com.codingapi.springboot.authorization.jdbc.AuthorizationJdbcDriver
-#spring.datasource.url=jdbc:mysql://localhost:3306/example?createDatabaseIfNotExist=true&useUnicode=true&characterEncoding=utf-8&useSSL=false&serverTimezone=Asia/Shanghai&allowPublicKeyRetrieval=true
-#spring.datasource.username=root
-#spring.datasource.password=lorne4j#2024
+#spring.datasource.url=jdbc:h2:file:./test.db
+#spring.jpa.database-platform=org.hibernate.dialect.H2Dialect
+#spring.jpa.hibernate.ddl-auto=create-drop
+#spring.jpa.show-sql=true
+
+spring.datasource.driver-class-name=com.codingapi.springboot.authorization.jdbc.AuthorizationJdbcDriver
+spring.datasource.url=jdbc:mysql://localhost:3306/example?createDatabaseIfNotExist=true&useUnicode=true&characterEncoding=utf-8&useSSL=false&serverTimezone=Asia/Shanghai&allowPublicKeyRetrieval=true
+spring.datasource.username=root
+spring.datasource.password=lorne4j#2024
 
 logging.level.com.codingapi.springboot.authorization=debug