上一篇 mybatis-plus sql改写插件_mybatis-plus sql插件-CSDN博客
去年基于租户过滤，权限拦截都是where tenant_id = XXX, 今年遇到一个物联网行业的需求，需要实现的是where sn in {sns} , 可能存在子查询，例如： where xinfo_id in (select x_id from x_info where sn in {sns}) ,这里的变量存在5个：xinfo_id,x_id,x_info,sn,sns 会有如下参数对象定义
java
@Data
class SubSelectParam{
String leftField; // xinfo_id
boolean isSubQuery; // true
String subTableName; // x_info
String selectField; // x_id
String subLeftField; // sn
ExpressionList stations; // sns
}以下是拦截完整代码
java
package com.sinexcel.framework.interceptor;
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.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.*;
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.*;
import java.util.stream.Collectors;
/**
* 数据权限处理器
* 根据用户角色和权限，拼接相应的 where 条件
* 支持select 子句、where 子句、join 子句、from 子句
*
* @since 3.4.1 +
*/
@Data
@NoArgsConstructor
@AllArgsConstructor
@SuppressWarnings({"rawtypes"})
public class DataPermissionInterceptor111 extends JsqlParserSupport implements InnerInterceptor {
public final Map<String, SubSelectParam> interceptorTables = new HashMap<String, SubSelectParam>(){{
// 不包含子查询的
put("table_name_nosub", new SubSelectParam("sn", Collections.emptyList()));
// 包含子查询的
put("table_name_sub", new SubSelectParam("sn", "other_table_name", "id", "sn",
Collections.emptyList()));
}};
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);
}
}
/**
* 递归分别识别 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].contains(".") ? fromItem[0].split("\\.")[1] : fromItem[0];
String aliasName = fromItem.length > 1 ? fromItem[1] : fromItem[0];
// 权限拦截器
SubSelectParam subSelectParam = interceptorTables.get(tableName.toLowerCase());
if(subSelectParam != null) {
InExpression inExpression;
if(subSelectParam.isSubQuery){
inExpression = new InExpression(new Column(aliasName + "." + subSelectParam.getLeftField()),
buildInSubQuery(subSelectParam));
}else{
inExpression = new InExpression(new Column(aliasName + "." + subSelectParam.getLeftField()),
subSelectParam.getStations());
}
sqlSegment = (sqlSegment == null ? inExpression : new AndExpression(sqlSegment, inExpression));
}else if("xxx".equalsIgnoreCase(tableName)){
// TODO 其他场景处理
}else{
// 无处理
}
if (null != sqlSegment) {
plainSelect.setWhere(sqlSegment);
}
}
public static SubSelect buildInSubQuery(SubSelectParam subSelectParam) {
SubSelect subSelectWrapper = new SubSelect()
.withSelectBody(
new PlainSelect()
.withSelectItems(List.of(new SelectExpressionItem(new Column(subSelectParam.getSelectField()))))
.withFromItem(new Table(subSelectParam.getSubTableName()))
.withWhere(new InExpression(new Column(subSelectParam.getSubLeftField()),
subSelectParam.getStations()))
);
return subSelectWrapper;
}
@Data
class SubSelectParam{
String leftField;
boolean isSubQuery;
String subTableName;
String selectField;
String subLeftField;
ExpressionList stations;
public SubSelectParam() {
}
/**
*
* @param leftField
* @param stations
*/
public SubSelectParam(String leftField, ExpressionList stations) {
this.leftField = leftField;
this.stations = stations;
this.isSubQuery = false;
}
/**
*
* @param leftField
* @param stations
*/
public SubSelectParam(String leftField, List<String> stations) {
super();
List<Expression> expressions = stations.stream().map(StringValue::new).collect(Collectors.toList());
this.stations = new ExpressionList(expressions);
this.leftField = leftField;
this.isSubQuery = false;
}
/**
*
* @param leftField
* @param subTableName
* @param selectField
* @param subLeftField
* @param stations
*/
public SubSelectParam(String leftField, String subTableName, String selectField, String subLeftField, ExpressionList stations) {
super();
this.leftField = leftField;
this.subTableName = subTableName;
this.selectField = selectField;
this.subLeftField = subLeftField;
this.stations = stations;
this.isSubQuery = true;
}
/**
*
* @param leftField
* @param subTableName
* @param selectField
* @param subLeftField
* @param stations
*/
public SubSelectParam(String leftField, String subTableName, String selectField, String subLeftField, List<String> stations) {
super();
List<Expression> expressions = stations.stream().map(StringValue::new).collect(Collectors.toList());
this.leftField = leftField;
this.subTableName = subTableName;
this.selectField = selectField;
this.subLeftField = subLeftField;
this.stations = new ExpressionList(expressions);
this.isSubQuery = true;
}
}
}