【SpringBoot整合mybatis-plus实现自定义拦截器-使用自定义注解】

自定义注解

package com.scm.common.datascope.annotation;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface CompanyDataScope {
    /**
     * 服务商别名- 对应业务表字段名称
     */
    public String serviceAlias() default "";

    /**
     * 客户别名- 对应业务表字段名称
     */
    public String customerAlias() default "";
}

第一步:实现MP自带拦截器 [InnerInterceptor]

package com.scm.common.datascope.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.inner.InnerInterceptor;
import lombok.*;
import net.sf.jsqlparser.expression.Expression;
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.SetOperationList;
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.List;


@Data
@NoArgsConstructor
@AllArgsConstructor
@ToString(callSuper = true)
@EqualsAndHashCode(callSuper = true)
public class CompanyDataInterceptor extends JsqlParserSupport implements InnerInterceptor {

    /**
     * 数据权限处理器
     */
    private CompanyDataHandler 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(this.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.setWhere((PlainSelect) selectBody, (String) obj);
        } else if (selectBody instanceof SetOperationList) {
            SetOperationList setOperationList = (SetOperationList) selectBody;
            List<SelectBody> selectBodyList = setOperationList.getSelects();
            selectBodyList.forEach(s -> this.setWhere((PlainSelect) s, (String) obj));
        }
    }

    /**
     * 设置 where 条件
     *
     * @param plainSelect  查询对象
     * @param whereSegment 查询条件片段
     */
    private void setWhere(PlainSelect plainSelect, String whereSegment) {

        Expression sqlSegment = this.dataPermissionHandler.getSqlSegment(plainSelect, whereSegment);
        if (null != sqlSegment) {
            plainSelect.setWhere(sqlSegment);
        }
    }
}

第二步: 拦截后数据处理方式

package com.scm.common.datascope.interceptor;

import cn.hutool.core.stream.CollectorUtil;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.scm.common.core.enums.CompanyData;
import com.scm.common.core.enums.DataScope;
import com.scm.common.datascope.annotation.CompanyDataScope;
import com.scm.common.security.utils.SecurityUtils;
import com.scm.system.api.model.LoginUser;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.Alias;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.HexValue;
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.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.PlainSelect;
import org.apache.commons.collections4.CollectionUtils;

import java.lang.reflect.Method;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;

@Slf4j
public class CompanyDataHandler {
    /**
     *
     * @param plainSelect  查询对象
     * @param whereSegment 查询条件片段
     * @return JSqlParser 条件表达式
     */
    @SneakyThrows(Exception.class)
    public Expression getSqlSegment(PlainSelect plainSelect, String whereSegment) {
        // 待执行 SQL Where 条件表达式
        Expression where = plainSelect.getWhere();
        if (where == null) {
            where = new HexValue(" 1 = 1 ");
        }
        log.info("开始进行权限过滤,where: {},mappedStatementId: {}", where, whereSegment);
        //获取mapper名称
        String className = whereSegment.substring(0, whereSegment.lastIndexOf("."));
        //获取方法名
        String methodName = whereSegment.substring(whereSegment.lastIndexOf(".") + 1);
        Table fromItem = (Table) plainSelect.getFromItem();
        // 有别名用别名，无别名用表名，防止字段冲突报错
        Alias fromItemAlias = fromItem.getAlias();
        String mainTableName = fromItemAlias == null ? fromItem.getName() : fromItemAlias.getName();
        //获取当前mapper 的方法
        Method[] methods = Class.forName(className).getMethods();
        //遍历判断mapper 的所以方法，判断方法上是否有 CompanyDataScope
        for (Method m : methods) {
            if (Objects.equals(m.getName(), methodName)) {
                CompanyDataScope annotation = m.getAnnotation(CompanyDataScope.class);
                if (annotation == null) {
                    return where;
                }
                // 1、当前用户
                LoginUser loginUser = SecurityUtils.getLoginUser();
                Set<String> dataScope= loginUser.getDataScope().stream().sorted().collect(Collectors.toSet());
                //如果数据权限为空则没有权限
                if ( dataScope.isEmpty()) return where;
                //2.获取用户对应企业数据权限
                Long companyId = loginUser.getSysUser().getScmCompanyId();
                //如果企业id为空则没有权限
                if (companyId == null) return where;
               // Set roleTypeSet =Sets.newHashSet("7");
                DataScope scopeType = CompanyData.getScope(dataScope);
                switch (scopeType) {
                    // 查看全部
                    case ALL:
                        return where;
                    case DEPT:
                        // 查看本服务商用户数据
                        // 创建IN 表达式
                        // 创建IN范围的元素集合
                        List<String> companyList = Lists.newArrayList(String.valueOf(companyId));
                        // 把集合转变为JSQLParser需要的元素列表
                        ItemsList deptList = new ExpressionList(companyList.stream().map(StringValue::new).collect(Collectors.toList()));
                        InExpression inExpressiondept = new InExpression(new Column(mainTableName + "." + annotation.serviceAlias()), deptList);
                        return new AndExpression(where, inExpressiondept);
                    case MYSELF:
                        // 查看客户的数据
                        //  = 表达式
                        EqualsTo usesEqualsTo = new EqualsTo();
                        usesEqualsTo.setLeftExpression(new Column(mainTableName + "." + annotation.customerAlias()));
                        usesEqualsTo.setRightExpression(new StringValue(String.valueOf(companyId)));
                        return new AndExpression(where, usesEqualsTo);
                    default:
                        break;
                }
            }

        }
        //说明无权查看，
        where = new HexValue(" 1 = 2 ");
        return where;
    }


}

第三步: 配置中心注入自定义拦截器

package com.scm.common.datascope.config;

import com.baomidou.mybatisplus.annotation.DbType;
import com.baomidou.mybatisplus.extension.plugins.MybatisPlusInterceptor;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import com.baomidou.mybatisplus.extension.plugins.inner.PaginationInnerInterceptor;
import com.scm.common.datascope.interceptor.CompanyDataHandler;
import com.scm.common.datascope.interceptor.CompanyDataInterceptor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.ArrayList;
import java.util.List;

@Configuration
public class MybatisConfig {
    /**
     * mybatis-plus分页插件配置
     *
     * @return PaginationInterceptor
     */
    @Bean(name = "mybatisPlusInterceptor")
    public MybatisPlusInterceptor mybatisPlusInterceptor() {
        MybatisPlusInterceptor interceptor = new MybatisPlusInterceptor();
        interceptor.addInnerInterceptor(new PaginationInnerInterceptor());
        return interceptor;
    }
	/**
	*注入自定义拦截器
	*/
    @Bean
    public CompanyDataInterceptor myInterceptor(MybatisPlusInterceptor mybatisPlusInterceptor) {
        CompanyDataInterceptor sql = new CompanyDataInterceptor();
        sql.setDataPermissionHandler(new CompanyDataHandler());
        List<InnerInterceptor> list = new ArrayList<>();
        // 添加数据权限插件
        list.add(sql);
        // 分页插件
        mybatisPlusInterceptor.setInterceptors(list);
        list.add(new PaginationInnerInterceptor(DbType.MYSQL));
        return sql;
    }
}

枚举的方式获取拦截权限

package com.scm.common.core.enums;

import lombok.AllArgsConstructor;
import lombok.Getter;

import java.util.Collection;

@AllArgsConstructor
@Getter
public enum CompanyData {

    // 枚举类型根据范围从前往后排列，避免影响getScope
    // Scope 数据权限范围 ： ALL（全部）、DEPT（服务商）、MYSELF（客户）
    DATA_MANAGER("运营", "6", DataScope.ALL),
    DATA_AUDITOR("服务商", "7",DataScope.DEPT),
    DATA_OPERATOR("客户", "8",DataScope.MYSELF);

    private String name;
    private String code;
    private DataScope scope;


    public static String getName(String code) {
        for (CompanyData type : CompanyData.values()) {
            if (type.getCode().equals(code)) {
                return type.getName();
            }
        }
        return null;
    }

    public static String getCode(String name) {
        for (CompanyData type : CompanyData.values()) {
            if (type.getName().equals(name)) {
                return type.getCode();
            }
        }
        return null;
    }

    public static DataScope getScope(Collection<String> code) {
        for (CompanyData type : CompanyData.values()) {
            for (String v : code) {
                if (type.getCode().equals(v)) {
                    return type.getScope();
                }
            }
        }
        return DataScope.MYSELF;
    }
}