add TableColumnAliasContext

Author: xlorne

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

@@ -23,26 +23,50 @@ private DataAuthorizationContext() {
         this.filters = new ArrayList<>();
     }
 
+    /**
+     * 添加数据权限过滤器
+     * @param filter 数据权限过滤器
+     */
     public void addDataAuthorizationFilter(DataAuthorizationFilter filter) {
         this.filters.add(filter);
     }
 
+    /**
+     *  清空数据权限过滤器
+     */
     public void clearDataAuthorizationFilters() {
         this.filters.clear();
     }
 
+    /**
+     * 列权限
+     * @param interceptState 拦截状态
+     * @param tableName 表名（或别名）
+     * @param columnName 列名 （或别名）
+     * @param value 值
+     * @return T
+     * @param <T> 泛型
+     */
     public <T> T columnAuthorization(SQLInterceptState interceptState, String tableName, String columnName, T value) {
         if (interceptState != null && interceptState.hasIntercept()) {
-//            String realTableName = interceptState.getTableName(tableName);
+            String realTableName = interceptState.getTableName(tableName);
+            String realColumnName = interceptState.getColumnName(tableName,columnName);
+
             for (DataAuthorizationFilter filter : filters) {
-                if (filter.supportColumnAuthorization(tableName, columnName, value)) {
-                    return filter.columnAuthorization(tableName, columnName, value);
+                if (filter.supportColumnAuthorization(realTableName, realColumnName, value)) {
+                    return filter.columnAuthorization(realTableName, realColumnName, value);
                 }
             }
         }
         return value;
     }
 
+    /**
+     * 行权限
+     * @param tableName 表名
+     * @param tableAlias 别名
+     * @return Condition 增加的过滤条件
+     */
     public Condition rowAuthorization(String tableName, String tableAlias) {
         if (StringUtils.hasText(tableName) && StringUtils.hasText(tableAlias)) {
             for (DataAuthorizationFilter filter : filters) {

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

@@ -3,20 +3,17 @@
 
 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.*;
+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 java.sql.SQLException;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
 
 /**
  * 数据权限 SQL 增强器
@@ -25,32 +22,31 @@ public class DataPermissionSQLEnhancer {
 
     private final String sql;
     private final RowHandler rowHandler;
+    private final TableColumnAliasHolder tableColumnAliasHolder;
+    private final Statement statement;
 
-    @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();
+    public DataPermissionSQLEnhancer(String sql, RowHandler rowHandler) throws SQLException {
+        try {
+            // 如何sql中存在? 则在?后面添加空格
+            this.sql = sql.replaceAll("\\?", " ? ");
+            this.rowHandler = rowHandler;
+            this.statement = CCJSqlParserUtil.parse(this.sql);
+            this.tableColumnAliasHolder = new TableColumnAliasHolder(statement);
+        } catch (Exception e) {
+            throw new SQLException(e);
+        }
     }
 
     // 获取增强后的SQL
     public String getNewSQL() throws SQLException {
         try {
-            Statement statement = CCJSqlParserUtil.parse(sql);
             if (statement instanceof Select) {
+                tableColumnAliasHolder.holderAlias();
                 Select select = (Select) statement;
                 PlainSelect plainSelect = select.getPlainSelect();
                 this.enhanceDataPermissionInSelect(plainSelect);
-                this.appendColumnAlias(plainSelect.getSelectItems());
-                aliasContext.print();
-                System.out.println(tableAlias);
                 return statement.toString();
             }
         } catch (Exception e) {
@@ -59,6 +55,10 @@ public String getNewSQL() throws SQLException {
         return sql;
     }
 
+    public TableColumnAliasContext getTableAlias() {
+        return tableColumnAliasHolder.getAliasContext();
+    }
+
 
     // 增强 SELECT 语句
     private void enhanceDataPermissionInSelect(PlainSelect plainSelect) throws Exception {
@@ -67,13 +67,11 @@ 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);
         }
@@ -83,65 +81,20 @@ 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();
         String aliaName = table.getAlias() != null ? table.getAlias().getName() : tableName;
-        tableAlias.put(aliaName, tableName);
         Condition condition = rowHandler.handler(plainSelect.toString(), tableName, aliaName);
         if (condition != null) {
             // 添加自定义条件

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

@@ -1,28 +1,62 @@
 package com.codingapi.springboot.authorization.enhancer;
 
-import java.util.HashMap;
+import lombok.Getter;
+
 import java.util.Map;
 
+/**
+ * 表字段别名
+ */
 public class TableColumnAlias {
 
-    private final String aliasName;
+    private final String parent;
+    @Getter
+    private final String tableName;
+    @Getter
     private final String columnName;
+    private final String aliasName;
 
-    private Map<String,TableColumnAlias> children;
+    public TableColumnAlias(String parent, String tableName, String columnName, String aliasName) {
+        this.parent = parent;
+        this.tableName = tableName;
+        this.columnName = columnName;
+        if (aliasName != null) {
+            this.aliasName = aliasName
+                    .replaceAll("`", "")
+                    .replaceAll("\"", "")
+                    .replaceAll("'", "")
+                    .trim();
+        } else {
+            this.aliasName = null;
+        }
+    }
 
-    public String getKey(){
-        return aliasName+"."+columnName;
+    /**
+     * 是否是表
+     *
+     * @param tableAlias 表别名
+     * @return 是否是表
+     */
+    public boolean isTable(Map<String, String> tableAlias) {
+        return tableAlias.containsKey(tableName);
     }
 
-    public TableColumnAlias(String aliasName, String columnName) {
-        this.aliasName = aliasName;
-        this.columnName = columnName;
-        this.children = new HashMap<>();
+
+    /**
+     * 获取父级key
+     *
+     * @return key
+     */
+    public String getParentKey() {
+        return parent + "." + aliasName;
     }
 
-    public TableColumnAlias add(String aliasName,String columnName){
-        TableColumnAlias tableColumnAlias = new TableColumnAlias(aliasName,columnName);
-        children.put(tableColumnAlias.getKey(),tableColumnAlias);
-        return tableColumnAlias;
+    /**
+     * 获取表别名key
+     *
+     * @return key
+     */
+    public String getTableAliasKey() {
+        return tableName + "." + aliasName;
     }
 }

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

@@ -1,27 +1,92 @@
 package com.codingapi.springboot.authorization.enhancer;
 
+import lombok.Getter;
+
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
+/**
+ *  表、字段别名上下文
+ */
 public class TableColumnAliasContext {
 
-    private final Map<String,TableColumnAlias> tableColumnAliasMap;
+    private final List<TableColumnAlias> columnAliases;
+    @Getter
+    private final Map<String, String> tableAlias;
+
+    private final Map<String, String> columnAliasMap;
+
+    protected TableColumnAliasContext() {
+        this.columnAliases = new ArrayList<>();
+        this.tableAlias = new HashMap<>();
+        this.columnAliasMap = new HashMap<>();
+    }
+
+    /**
+     * 添加表别名
+     * @param tableAlias 表别名
+     * @param tableName 表名
+     */
+    protected void addTable(String tableAlias, String tableName) {
+        this.tableAlias.put(tableAlias, tableName);
+    }
 
-    public TableColumnAliasContext() {
-        this.tableColumnAliasMap = new HashMap<>();
+    /**
+     * 添加字段别名
+     * @param parent 父级（上级别名）
+     * @param tableName 表名
+     * @param columnName 字段名
+     * @param aliasName 别名
+     */
+    protected void addColumn(String parent, String tableName, String columnName, String aliasName) {
+        TableColumnAlias tableColumnAlias = new TableColumnAlias(parent, tableName, columnName, aliasName);
+        columnAliases.add(tableColumnAlias);
     }
 
-    public TableColumnAlias add(String aliasName,String columnName){
-        TableColumnAlias tableColumnAlias = new TableColumnAlias(aliasName,columnName);
-        tableColumnAliasMap.put(tableColumnAlias.getKey(),tableColumnAlias);
-        return tableColumnAlias;
+    /**
+     * 列别名转换为map
+     */
+    protected void columnKeyToMap() {
+        for (TableColumnAlias tableColumnAlias : columnAliases) {
+            if (tableColumnAlias.isTable(tableAlias)) {
+                String parentKey = tableColumnAlias.getParentKey();
+                String tableAliasName = tableColumnAlias.getTableName();
+                String tableName = this.getTableName(tableAliasName);
+                String columnValue = tableName + "." + tableColumnAlias.getColumnName();
+                columnAliasMap.put(parentKey, columnValue);
+                columnAliasMap.put(tableColumnAlias.getTableAliasKey(), columnValue);
+            }
+        }
+    }
+
+    /**
+     * 获取表名（真实表名）
+     * @param tableName 表名或表别名
+     * @return 真实表名
+     */
+    public String getTableName(String tableName) {
+        String value = tableAlias.get(tableName);
+        if (value != null) {
+            return value;
+        }
+        return tableName;
     }
 
 
-    public void print(){
-        for(String key:tableColumnAliasMap.keySet()){
-            TableColumnAlias tableColumnAlias = tableColumnAliasMap.get(key);
-            System.out.println(tableColumnAlias.getKey());
+    /**
+     * 获取字段名（真实字段名）
+     * @param tableName 表名或表别名
+     * @param columnName 字段名
+     * @return 真实字段名
+     */
+    public String getColumnName(String tableName, String columnName) {
+        String key = tableName + "." + columnName;
+        String value = columnAliasMap.get(key);
+        if (value != null) {
+            return value.split("\\.")[1];
         }
+        return columnName;
     }
 }