新增租户插件

src/main/java/com/baomidou/mybatisplus/plugins/TenancyInterceptor.java

@@ -18,10 +18,12 @@ package com.baomidou.mybatisplus.plugins;
 import java.util.Properties;
 
 import org.apache.ibatis.executor.Executor;
+import org.apache.ibatis.executor.statement.StatementHandler;
 import org.apache.ibatis.mapping.MappedStatement;
 import org.apache.ibatis.plugin.Interceptor;
 import org.apache.ibatis.plugin.Intercepts;
 import org.apache.ibatis.plugin.Invocation;
+import org.apache.ibatis.plugin.Plugin;
 import org.apache.ibatis.plugin.Signature;
 import org.apache.ibatis.session.ResultHandler;
 import org.apache.ibatis.session.RowBounds;
@@ -32,7 +34,7 @@ import org.apache.ibatis.session.RowBounds;
  * </p>
  *
  * @author hubin
- * @Date 2016-08-16
+ * @since 2016-08-16
  */
 @Intercepts({
         @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
@@ -42,16 +44,20 @@ public class TenancyInterceptor implements Interceptor {
 
     @Override
     public Object intercept(Invocation invocation) throws Throwable {
+
         return null;
     }
 
     @Override
     public Object plugin(Object target) {
-        return null;
+        if (target instanceof StatementHandler) {
+            return Plugin.wrap(target, this);
+        }
+        return target;
     }
 
     @Override
     public void setProperties(Properties properties) {
-
+        // to do nothing
     }
 }

src/main/java/com/baomidou/mybatisplus/plugins/pagination/optimize/AliDruidCountOptimize.java

@@ -25,7 +25,7 @@ import com.baomidou.mybatisplus.parser.SqlInfo;
  * </p>
  *
  * @author hubin
- * @Date 2017-06-20
+ * @since 2017-06-20
  */
 public class AliDruidCountOptimize extends AbstractSqlParser {
 

src/main/java/com/baomidou/mybatisplus/plugins/pagination/optimize/DefaultCountOptimize.java

@@ -24,7 +24,7 @@ import com.baomidou.mybatisplus.parser.SqlInfo;
  * </p>
  *
  * @author hubin
- * @Date 2017-06-20
+ * @since 2017-06-20
  */
 public class DefaultCountOptimize extends AbstractSqlParser {
 

src/main/java/com/baomidou/mybatisplus/plugins/pagination/optimize/JsqlParserCountOptimize.java

@@ -41,7 +41,7 @@ import net.sf.jsqlparser.statement.select.SelectItem;
  * </p>
  *
  * @author hubin
- * @Date 2017-06-20
+ * @since 2017-06-20
  */
 public class JsqlParserCountOptimize extends AbstractSqlParser {
     private static final List<SelectItem> countSelectItem = countSelectItem();

src/main/java/com/baomidou/mybatisplus/plugins/tenancy/TenancySqlParser.java

@@ -1,16 +1,64 @@
+/**
+ * Copyright (c) 2011-2020, hubin (jobob@qq.com).
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
 package com.baomidou.mybatisplus.plugins.tenancy;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import com.baomidou.mybatisplus.parser.AbstractSqlParser;
 import com.baomidou.mybatisplus.parser.SqlInfo;
 
+import net.sf.jsqlparser.JSQLParserException;
+import net.sf.jsqlparser.expression.BinaryExpression;
+import net.sf.jsqlparser.expression.Expression;
+import net.sf.jsqlparser.expression.StringValue;
+import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
+import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
+import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
+import net.sf.jsqlparser.expression.operators.relational.InExpression;
+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.delete.Delete;
 import net.sf.jsqlparser.statement.insert.Insert;
+import net.sf.jsqlparser.statement.select.FromItem;
+import net.sf.jsqlparser.statement.select.Join;
+import net.sf.jsqlparser.statement.select.LateralSubSelect;
+import net.sf.jsqlparser.statement.select.PlainSelect;
+import net.sf.jsqlparser.statement.select.Select;
 import net.sf.jsqlparser.statement.select.SelectBody;
+import net.sf.jsqlparser.statement.select.SelectExpressionItem;
+import net.sf.jsqlparser.statement.select.SetOperationList;
+import net.sf.jsqlparser.statement.select.SubJoin;
+import net.sf.jsqlparser.statement.select.SubSelect;
+import net.sf.jsqlparser.statement.select.ValuesList;
+import net.sf.jsqlparser.statement.select.WithItem;
 import net.sf.jsqlparser.statement.update.Update;
 
 /**
- * Created by jobob on 17/6/20.
+ * <p>
+ * 租户 SQL 解析
+ * </p>
+ *
+ * @author hubin
+ * @since 2017-06-20
  */
 public class TenancySqlParser extends AbstractSqlParser {
+    private TenantInfo tenantInfo;
 
     public TenancySqlParser(String sql, String dbType) {
         super(sql, dbType);
@@ -18,40 +66,247 @@ public class TenancySqlParser extends AbstractSqlParser {
 
     @Override
     public SqlInfo optimizeSql() {
-        return null;
+        String sql = this.getSql();
+        //logger.debug("old sql:{}", sql);
+        Statement stmt = null;
+        try {
+            stmt = CCJSqlParserUtil.parse(sql);
+        } catch (JSQLParserException e) {
+            //logger.debug("解析", e);
+            //logger.error("解析sql[{}]失败\n原因:{}", sql, e.getMessage());
+            //如果解析失败不进行任何处理防止业务中断
+            return null;
+        }
+        if (stmt instanceof Insert) {
+            processInsert((Insert) stmt);
+        } else if (stmt instanceof Select) {
+            processSelectBody(((Select) stmt).getSelectBody());
+        } else if (stmt instanceof Update) {
+            processUpdate((Update) stmt);
+        }
+        //logger.debug("new sql:{}", stmt);
+        SqlInfo sqlInfo = SqlInfo.newInstance();
+        sqlInfo.setSql(stmt.toString());
+        return sqlInfo;
     }
 
     /**
      * <p>
      * select 语句处理
      * </p>
-     *
-     * @param selectBody
      */
     public void processSelectBody(SelectBody selectBody) {
-
+        if (selectBody instanceof PlainSelect) {
+            processPlainSelect((PlainSelect) selectBody);
+        } else if (selectBody instanceof WithItem) {
+            WithItem withItem = (WithItem) selectBody;
+            if (withItem.getSelectBody() != null) {
+                processSelectBody(withItem.getSelectBody());
+            }
+        } else {
+            SetOperationList operationList = (SetOperationList) selectBody;
+            if (operationList.getSelects() != null && operationList.getSelects().size() > 0) {
+                List<SelectBody> plainSelects = operationList.getSelects();
+                for (SelectBody plainSelect : plainSelects) {
+                    processSelectBody(plainSelect);
+                }
+            }
+        }
     }
 
     /**
      * <p>
      * insert 语句处理
      * </p>
-     *
-     * @param insert
      */
     public void processInsert(Insert insert) {
-
+        if (doTableFilter(
+                insert.getTable().getName()
+        )) {
+            insert.getColumns().add(new Column(this.tenantInfo.getTenantIdColumn()));
+            if (insert.getSelect() != null) {
+                processPlainSelect((PlainSelect) insert.getSelect().getSelectBody(), true);
+            } else if (insert.getItemsList() != null) {
+                ((ExpressionList) insert.getItemsList()).getExpressions().add(new StringValue("," + this.tenantInfo.getTenantId() + ","));
+            } else {
+                //
+                throw new RuntimeException("无法处理的 sql");
+            }
+        }
     }
 
     /**
      * <p>
      * update 语句处理
      * </p>
-     *
-     * @param update
      */
     public void processUpdate(Update update) {
+        //获得where条件表达式
+        Expression where = update.getWhere();
+        EqualsTo equalsTo = new EqualsTo();
+        if (where instanceof BinaryExpression) {
+            equalsTo.setLeftExpression(new Column(this.tenantInfo.getTenantIdColumn()));
+            equalsTo.setRightExpression(new StringValue("," + tenantInfo.getTenantId() + ","));
+            AndExpression andExpression = new AndExpression(equalsTo, where);
+            update.setWhere(andExpression);
+        } else {
+            equalsTo.setLeftExpression(new Column(this.tenantInfo.getTenantIdColumn()));
+            equalsTo.setRightExpression(new StringValue("," + tenantInfo.getTenantId() + ","));
+            update.setWhere(equalsTo);
+        }
+    }
+
+
+    /**
+     * <p>
+     * delete 语句处理
+     * </p>
+     */
+    public void processDelete(Delete delete) {
 
     }
 
+    /**
+     * 处理PlainSelect
+     */
+    public void processPlainSelect(PlainSelect plainSelect) {
+        processPlainSelect(plainSelect, false);
+    }
+
+    /**
+     * 处理PlainSelect
+     *
+     * @param plainSelect
+     * @param addColumn   是否添加租户列,insert into select语句中需要
+     */
+
+    public void processPlainSelect(PlainSelect plainSelect, boolean addColumn) {
+        FromItem fromItem = plainSelect.getFromItem();
+        if (fromItem instanceof Table) {
+            Table fromTable = (Table) fromItem;
+            if (doTableFilter(fromTable.getName())) {
+                plainSelect.setWhere(builderExpression(plainSelect.getWhere(), fromTable));
+                if (addColumn)
+                    plainSelect.getSelectItems().add(new SelectExpressionItem(new Column("'" + this.tenantInfo.getTenantId() + "'")));
+            }
+        } else {
+            processFromItem(fromItem);
+        }
+        List<Join> joins = plainSelect.getJoins();
+        if (joins != null && joins.size() > 0) {
+            for (Join join : joins) {
+                processJoin(join);
+                processFromItem(join.getRightItem());
+            }
+        }
+    }
+
+    /**
+     * 处理子查询等
+     *
+     * @param fromItem
+     */
+    public void processFromItem(FromItem fromItem) {
+        if (fromItem instanceof SubJoin) {
+            SubJoin subJoin = (SubJoin) fromItem;
+            if (subJoin.getJoin() != null) {
+                processJoin(subJoin.getJoin());
+            }
+            if (subJoin.getLeft() != null) {
+                processFromItem(subJoin.getLeft());
+            }
+        } else if (fromItem instanceof SubSelect) {
+            SubSelect subSelect = (SubSelect) fromItem;
+            if (subSelect.getSelectBody() != null) {
+                processSelectBody(subSelect.getSelectBody());
+            }
+        } else if (fromItem instanceof ValuesList) {
+
+        } else if (fromItem instanceof LateralSubSelect) {
+            LateralSubSelect lateralSubSelect = (LateralSubSelect) fromItem;
+            if (lateralSubSelect.getSubSelect() != null) {
+                SubSelect subSelect = lateralSubSelect.getSubSelect();
+                if (subSelect.getSelectBody() != null) {
+                    processSelectBody(subSelect.getSelectBody());
+                }
+            }
+        }
+    }
+
+    /**
+     * 处理联接语句
+     *
+     * @param join
+     */
+    public void processJoin(Join join) {
+        if (join.getRightItem() instanceof Table) {
+            Table fromTable = (Table) join.getRightItem();
+            if (doTableFilter(fromTable.getName())) {
+                join.setOnExpression(builderExpression(join.getOnExpression(), fromTable));
+            }
+
+        }
+    }
+
+
+    /**
+     * 处理条件
+     * TODO 未解决sql注入问题(考虑替换StringValue为LongValue),因为线上数据库租户字段为int暂时不存在注入问题
+     *
+     * @param expression
+     * @param table
+     * @return
+     */
+    public Expression builderExpression(Expression expression, Table table) {
+        Expression tenantExpression = null;
+        String[] tenantIds = this.tenantInfo.getTenantId().split(",");
+        //当传入table时,字段前加上别名或者table名
+        //别名优先使用
+        StringBuilder tenantIdColumnName = new StringBuilder();
+        if (table != null) {
+            tenantIdColumnName.append(table.getAlias() != null ? table.getAlias().getName() : table.getName());
+            tenantIdColumnName.append(".");
+        }
+        tenantIdColumnName.append(this.tenantInfo.getTenantIdColumn());
+        //生成字段名
+        Column tenantColumn = new Column(tenantIdColumnName.toString());
+
+        if (tenantIds.length == 1) {
+            EqualsTo equalsTo = new EqualsTo();
+            tenantExpression = equalsTo;
+            equalsTo.setLeftExpression(tenantColumn);
+            equalsTo.setRightExpression(new StringValue("'" + tenantIds[0] + "'"));
+        } else {
+            //多租户身份
+            InExpression inExpression = new InExpression();
+            tenantExpression = inExpression;
+            inExpression.setLeftExpression(tenantColumn);
+            List<Expression> valueList = new ArrayList<>();
+            for (String tid : tenantIds) {
+                valueList.add(new StringValue("'" + tid + "'"));
+            }
+            inExpression.setRightItemsList(new ExpressionList(valueList));
+        }
+
+        //加入判断防止条件为空时生成 "and null" 导致查询结果为空
+        if (expression == null) {
+            return tenantExpression;
+        } else {
+            if (expression instanceof BinaryExpression) {
+                BinaryExpression binaryExpression = (BinaryExpression) expression;
+                if (binaryExpression.getLeftExpression() instanceof FromItem) {
+                    processFromItem((FromItem) binaryExpression.getLeftExpression());
+                }
+                if (binaryExpression.getRightExpression() instanceof FromItem) {
+                    processFromItem((FromItem) binaryExpression.getRightExpression());
+                }
+            }
+            return new AndExpression(tenantExpression, expression);
+        }
+
+    }
+
+    private boolean doTableFilter(String table) {
+        return true;
+    }
 }

src/main/java/com/baomidou/mybatisplus/plugins/tenancy/TenantInfo.java

@@ -0,0 +1,38 @@
+/**
+ * Copyright (c) 2011-2020, hubin (jobob@qq.com).
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.baomidou.mybatisplus.plugins.tenancy;
+
+import java.util.Properties;
+
+/**
+ * <p>
+ * 租户信息
+ * </p>
+ *
+ * @author hubin
+ * @since 2017-06-20
+ */
+public interface TenantInfo {
+
+    TenantInfo setFilterConfig(Properties properties);
+
+    String getTenantId();
+
+    String getTenantIdColumn();
+
+    TenantInfo setTenantIdColumn(String tenantIdColumn);
+
+}

src/main/java/com/baomidou/mybatisplus/plugins/tenancy/handler/RegxTenancyHandler.java

@@ -0,0 +1,125 @@
+/**
+ * Copyright (c) 2011-2020, hubin (jobob@qq.com).
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.baomidou.mybatisplus.plugins.tenancy.handler;
+
+import java.util.Properties;
+import java.util.regex.Pattern;
+
+import com.baomidou.mybatisplus.toolkit.PluginUtils;
+
+/**
+ * <p>
+ * 租户信息正则处理器
+ * </p>
+ *
+ * @author hubin
+ * @since 2017-06-20
+ */
+public class RegxTenancyHandler implements TenancyHandler {
+    //默认过滤还是忽略 true表示按租户过滤
+    private boolean filterDefault = false;
+    private Pattern tablePatterns[];
+    private Pattern statementPatterns[];
+
+    @Override
+    public void setConfig(Properties properties) {
+        this.setFilterStatementRegexStr(PluginUtils.getProperty(properties, "filterStatementRegexStr"));
+        this.setFilterTableRegexStr(PluginUtils.getProperty(properties, "filterTableRegexStr"));
+        String filterDefault = PluginUtils.getProperty(properties, "filterDefault");
+        if (filterDefault != null) this.setFilterDefault(Boolean.valueOf(filterDefault));
+    }
+
+    @Override
+    public boolean doTable(String table) {
+        boolean isOk = filterDefault;
+        if (this.tablePatterns != null) {
+            for (Pattern p : this.tablePatterns) {
+                if (p.matcher(table).find()) {
+                    isOk = !filterDefault;
+                    break;
+                }
+            }
+        }
+        // logger.debug("table:{}  isOK:{}   tablePatterns:{}",table,isOk,tablePatterns);
+        return isOk;
+    }
+
+    @Override
+    public boolean doStatement(String statementId) {
+        boolean isOk = filterDefault;
+        if (this.statementPatterns != null) {
+            for (Pattern p : this.statementPatterns) {
+                if (p.matcher(statementId).find()) {
+                    isOk = !filterDefault;
+                    break;
+                }
+            }
+        }
+        //logger.debug("statementId:{}  isOK:{}   statementPatterns:{}",statementId,isOk,statementPatterns);
+        return isOk;
+    }
+
+    public static Pattern[] compile(String patterString) {
+        if (patterString == null) return new Pattern[]{};
+        String[] patterStrings = patterString.split(",");
+        return compile(patterStrings);
+    }
+
+    public static Pattern[] compile(String[] patterStrings) {
+        if (patterStrings == null) return new Pattern[]{};
+        Pattern[] patterns = new Pattern[patterStrings.length];
+        for (int i = 0; i < patterStrings.length; i++) {
+            Pattern pattern = Pattern.compile(patterStrings[i]);
+            patterns[i] = pattern;
+        }
+        return patterns;
+    }
+
+    public RegxTenancyHandler setFilterTableRegexStr(String tableRegexStr) {
+        if (tableRegexStr == null) {
+            return this;
+        }
+        this.tablePatterns = compile(tableRegexStr);
+        return this;
+    }
+
+    public RegxTenancyHandler setFilterTableRegexArr(String[] tableRegexArr) {
+        if (tableRegexArr == null) return this;
+        this.tablePatterns = compile(tableRegexArr);
+        return this;
+    }
+
+    public RegxTenancyHandler setFilterStatementRegexStr(String statementRegexStr) {
+        if (statementRegexStr == null) {
+            return this;
+        }
+        this.statementPatterns = compile(statementRegexStr);
+        return this;
+    }
+
+    public RegxTenancyHandler setFilterStatementRegexArr(String[] statementRegexArr) {
+        if (statementRegexArr == null) {
+            return this;
+        }
+        this.statementPatterns = compile(statementRegexArr);
+        return this;
+    }
+
+    public RegxTenancyHandler setFilterDefault(boolean filterDefault) {
+        this.filterDefault = filterDefault;
+        return this;
+    }
+}

src/main/java/com/baomidou/mybatisplus/plugins/tenancy/handler/TenancyHandler.java

@@ -0,0 +1,60 @@
+/**
+ * Copyright (c) 2011-2020, hubin (jobob@qq.com).
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.baomidou.mybatisplus.plugins.tenancy.handler;
+
+import java.util.Properties;
+
+/**
+ * <p>
+ * 租户信息处理器
+ * </p>
+ *
+ * @author hubin
+ * @since 2017-06-20
+ */
+public interface TenancyHandler {
+
+    /**
+     * <p>
+     * 配置设置
+     * </p>
+     *
+     * @param properties mybatis Interceptor setProperties
+     */
+    void setConfig(Properties properties);
+
+    /**
+     * <p>
+     * 按照表名处理
+     * </p>
+     *
+     * @param table 表名
+     * @return true 执行，false 不执行
+     */
+    boolean doTable(String table);
+
+
+    /**
+     * <p>
+     * 按照statementId处理
+     * </p>
+     *
+     * @param statementId mybatis statementId
+     * @return true 执行，false 不执行
+     */
+    boolean doStatement(String statementId);
+
+}

src/main/java/com/baomidou/mybatisplus/plugins/tenancy/handler/package-info.java

@@ -0,0 +1,4 @@
+/**
+ * 租户信息处理相关类
+ */
+package com.baomidou.mybatisplus.plugins.tenancy.handler;

src/main/java/com/baomidou/mybatisplus/plugins/tenancy/package-info.java

@@ -0,0 +1,4 @@
+/**
+ * mybatis 租户插件实现类
+ */
+package com.baomidou.mybatisplus.plugins.tenancy;

src/main/java/com/baomidou/mybatisplus/toolkit/PluginUtils.java

@@ -1,22 +1,44 @@
+/**
+ * Copyright (c) 2011-2020, hubin (jobob@qq.com).
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
 package com.baomidou.mybatisplus.toolkit;
 
 import java.lang.reflect.Proxy;
+import java.util.Properties;
 
 import org.apache.ibatis.reflection.MetaObject;
 import org.apache.ibatis.reflection.SystemMetaObject;
 
 /**
+ * <p>
  * 插件工具类
+ * </p>
  *
- * @author TaoYu
+ * @author TaoYu , hubin
+ * @since 2017-06-20
  */
 public final class PluginUtils {
 
     private PluginUtils() {
+        // to do nothing
     }
 
     /**
+     * <p>
      * 获得真正的处理对象,可能多层代理.
+     * </p>
      */
     public static Object realTarget(Object target) {
         if (Proxy.isProxyClass(target.getClass())) {
@@ -26,4 +48,13 @@ public final class PluginUtils {
         return target;
     }
 
+    /**
+     * <p>
+     * 根据 key 获取 Properties 的值
+     * </p>
+     */
+    public static String getProperty(Properties properties, String key) {
+        String value = properties.getProperty(key);
+        return StringUtils.isEmpty(value) ? null : value;
+    }
 }