mybatis-plus sql改写插件

import com.baomidou.mybatisplus.core.plugins.InterceptorIgnoreHelper;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport;
import com.baomidou.mybatisplus.extension.plugins.handler.DataPermissionHandler;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
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.schema.Column;
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.SubSelect;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.sql.SQLException;
import java.util.Arrays;
import java.util.List;

/**
 * 数据权限处理器
 * 根据用户角色和权限，拼接相应的 where 条件
 * 支持select 子句、where 子句、join 子句、from 子句、with临时表
 *
 * @since 3.4.1 +
 */
@Data
@NoArgsConstructor
@AllArgsConstructor
@SuppressWarnings({"rawtypes"})
public class DataPermissionInterceptor1 extends JsqlParserSupport implements InnerInterceptor {
	public static final List<String> interceptorTableNames = Arrays.asList("order","order_detail");
    private DataPermissionHandler dataPermissionHandler;

    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        if (InterceptorIgnoreHelper.willIgnoreDataPermission(ms.getId())) return;
        PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
        mpBs.sql(parserSingle(mpBs.sql(), ms.getId()));
    }


    @Override
    protected void processSelect(Select select, int index, String sql, Object obj) {
        SelectBody selectBody = select.getSelectBody();
        // 普通查询处理分支
        if (selectBody instanceof PlainSelect) {
            this.recursiveSubSqlSetWhere((PlainSelect) selectBody, (String) obj);
        }
        // with as 子句处理分支
        if (select.getWithItemsList() != null && select.getWithItemsList().size() > 0) {
            select.getWithItemsList().forEach(wi -> {
                if (wi.getSelectBody() instanceof PlainSelect) {
                    this.recursiveSubSqlSetWhere((PlainSelect) wi.getSelectBody(), (String) obj);
                }
            });
        }
    }

    /**
     * 递归分别识别 select 子句中的子查询
     *            from 子句中可能存在子查询
     *            where 子句中可能存在子查询
     *            join 子句中可能存在子查询
     * @param plainSelect
     * @param whereSegment
     */
    private void recursiveSubSqlSetWhere(PlainSelect plainSelect, String whereSegment) {
        this.setWhere(plainSelect, whereSegment);
        // from 子句中可能存在子查询
        if(plainSelect.getFromItem() instanceof SubSelect){
            SubSelect subSelect = (SubSelect)plainSelect.getFromItem();
            subSelectSetWhere(subSelect, whereSegment);
        }
        // join 子句中可能存在子查询
        if (plainSelect.getJoins() != null) {
            plainSelect.getJoins().forEach(j -> {
                if (j.getRightItem() instanceof SubSelect) {
                    SubSelect subSelect = (SubSelect)j.getRightItem();
                    subSelectSetWhere(subSelect, whereSegment);
                }
            });
        }
        // select 子句中可能存在子查询
        if (plainSelect.getSelectItems() != null) {
            plainSelect.getSelectItems().forEach(s -> {
                if (s instanceof SelectExpressionItem) {
                    Expression expression = ((SelectExpressionItem) s).getExpression();
                    if(expression instanceof SubSelect){
                        SubSelect subSelect = (SubSelect)expression;
                        subSelectSetWhere(subSelect, whereSegment);
                    }
                }
            });
        }
        // where 子句中可能存在子查询
        if (plainSelect.getWhere() != null) {
            if(plainSelect.getWhere() instanceof BinaryExpression){
                BinaryExpression expression = (BinaryExpression) plainSelect.getWhere();
                recursiveWhereExpression(expression, whereSegment);
            }
        }
        // having 子句中可能存在子查询
        if (plainSelect.getHaving() != null) {
            if(plainSelect.getHaving() instanceof BinaryExpression){
                BinaryExpression expression = (BinaryExpression) plainSelect.getWhere();
                recursiveWhereExpression(expression, whereSegment);
            }
        }
    }

    /**
     * 递归找出where 子句中的子查询，并设置租户条件
     * @param expression
     * @param whereSegment
     */
    private void recursiveWhereExpression(BinaryExpression expression, String whereSegment){
        if(expression.getLeftExpression() instanceof BinaryExpression){
            recursiveWhereExpression((BinaryExpression)expression.getLeftExpression(), whereSegment);
        }
        if(expression.getRightExpression() instanceof BinaryExpression){
            recursiveWhereExpression((BinaryExpression)expression.getRightExpression(), whereSegment);
        }
        if(expression.getLeftExpression() instanceof SubSelect){
            SubSelect subSelect = (SubSelect)expression.getLeftExpression();
            subSelectSetWhere(subSelect, whereSegment);
        }
        if(expression.getRightExpression() instanceof SubSelect){
            SubSelect subSelect = (SubSelect)expression.getRightExpression();
            subSelectSetWhere(subSelect, whereSegment);
        }
    }

    /**
     * 子句递归设置 where 条件
     * @param subSelect
     * @param whereSegment
     */
    private void subSelectSetWhere(SubSelect subSelect, String whereSegment) {
        SelectBody selectBody = subSelect.getSelectBody();
        if (selectBody instanceof PlainSelect) {
            this.recursiveSubSqlSetWhere((PlainSelect) selectBody, whereSegment);
        }
    }

    /**
     * 设置 where 条件
     *
     * @param plainSelect  查询对象
     * @param whereSegment 查询条件片段
     */
    protected void setWhere(PlainSelect plainSelect, String whereSegment) {
    	Expression sqlSegment = plainSelect.getWhere();
	
        String[] fromItem = plainSelect.getFromItem().toString().replaceAll("`", "").split(" ");
        String tableName = fromItem[0].indexOf(".") > -1 ? fromItem[0].split("\\.")[1] : fromItem[0];
        String aliasName = fromItem.length > 1 ? fromItem[1] : fromItem[0];

        // 权限拦截器
        if(interceptorTableNames.contains(tableName.toLowerCase())) {
            EqualsTo equalsTo = new EqualsTo();
            equalsTo.setLeftExpression(new Column(aliasName + ".user_id"));
            equalsTo.setRightExpression(new StringValue("userid"));// 替换成实际的值
            sqlSegment = (sqlSegment == null ? equalsTo : new AndExpression(sqlSegment, equalsTo));
        }


        if (null != sqlSegment) {
            plainSelect.setWhere(sqlSegment);
        }
    }
}