MybatisPlus实现数据权限

前言

在管理系统中，我们经常要做数据权限限制，当前登录用户可以看哪一部分数据（一般为当前归属部门及其下级部门），但是每次做这样的功能都需要熟悉整个系统的权限规划，还要做很多重复性工作，是否可以将此功能封装起来呢？

实现思路

使用MybatisPlus的拦截器，拦截每一个查询方法，并根据执行方法判断是否包含自定义的注解，
如果包含则将当前执行的SQL添加上自定义的数据权限筛选条件
步骤如下

  1. 创建自定义注解
  2. 创建自定义拦截器
  3. 创建通用处理器接口（支持自定义处理逻辑）
  4. 将拦截器注册到MybatisPlus
  5. 添加工具类配合MybatisPlus的默认Wrapper使用

创建注解

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface DataScope {
    /**
     * <p>
     *     数据归属机构ID字段名
     * </p>
     */
    String orgIdColumnName() default "org_id";

    /**
     * 数据创建人ID字段名
     */
    String createByColumnName() default "create_by";

    /**
     * 处理器
     */
    Class<? extends DataPermissionHandler> handler() default DefaultDataPermissionHandler.class;
}

创建拦截器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
@Slf4j
public class DataPermissionInterceptor extends JsqlParserSupport implements InnerInterceptor {
    /**
     * 方法
     */
    private static final Map<String, DataScope> DATA_SCOPE_CACHE = new ConcurrentHashMap<>();
    private static final Map<String, Method> DATA_SCOPE_METHOD_CACHE = new ConcurrentHashMap<>();
    /**
     *  查找过一次，且不包含注解的方法缓存
     */
    private static final Set<String> IGNORE_METHOD_CACHE = new ConcurrentHashSet<>();


    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
        // 没有注解不进行操作
        DataScope dataScope = getDataScope(ms.getId());
        if (dataScope == null) {
            return;
        }
        // 调用 processSelect 使用 jsql处理SQL
        PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
        mpBs.sql(parserSingle(mpBs.sql(), ms.getId()));
    }
    // 真实处理逻辑
    @Override
    protected void processSelect(Select select, int index, Object obj) {
        PlainSelect plainSelect = (PlainSelect) select.getSelectBody();
        Expression where = plainSelect.getWhere();
        try {
            // 调用handler的方法名
            final String methodName = "getSqlSegment";
            // 调用handler的参数类型
            final Class[] methodParmeters = {Expression.class, String.class, DataScope.class, Method.class};
            // 此处obj为beforeQuery方法中传入的ms.getId()，即为 mappedStatementId
            final String mappedStatementId = Func.toStr(obj);
            // 获取注解，如果可以走到这里说明注解一定不为空
            DataScope dataScope = getDataScope(mappedStatementId);
            // 获取处理器并调用处理方法
            Class<? extends DataPermissionHandler> handler = Objects.requireNonNull(dataScope).handler();
            // 获取处理方法，强制与 DataPermissionHandler 方法对应
            Method getSqlSegmentMethod = ClassUtil.getDeclaredMethod(handler, methodName, methodParmeters);
            Assert.notNull(getSqlSegmentMethod, "[%s]未找到处理方法[%s]", handler.getName(), methodName);
            // 当前执行的mapper方法，方法在获取 getDataScope 获取注解成功后缓存的
            Method method = DATA_SCOPE_METHOD_CACHE.get(mappedStatementId);
            Object sqlSegment = ReflectUtil.invoke(Singleton.get(handler), getSqlSegmentMethod, where, mappedStatementId, dataScope, method);
            if (Func.isNotEmpty(sqlSegment)) {
                plainSelect.setWhere((Expression) sqlSegment);
            }
        } catch (Exception e) {
            // 调用处理方法出现任何错误时，创建一个不可能实现的Where条件表达式，使其不能查出任何数据，并抛出异常
            EqualsTo equalsTo = new EqualsTo();
            equalsTo.setLeftExpression(new LongValue(1));
            equalsTo.setRightExpression(new LongValue(2));
            plainSelect.setWhere(equalsTo);
            throw e;
        }
    }
    // 获取当前class与其所有实现的接口上的所有方法（由于Mybatis的Mapper都是接口）
    private void addAllMethod(List<Class> clazzs, List<Method> allMethods) {
        if (clazzs == null) {
            return;
        }
        clazzs.forEach(c -> {
            if (c != null) {
                allMethods.addAll(Arrays.asList(c.getMethods()));
                Class[] interfaces = c.getInterfaces();
                addAllMethod(Arrays.asList(interfaces), allMethods);
            }
        });
    }


    /**
     * 根据mappedStatementId获取方法上的注解,并加入缓存
     *
     * @param mappedStatementId mybatis的mappedStatementId 示例 top.kthirty.mapper.UserMapper.list
     * @return 获取注解，找不到时返回Null
     */
    private DataScope getDataScope(String mappedStatementId) {
        // 查找过一次，未发现注解，不再查找
        if(IGNORE_METHOD_CACHE.contains(mappedStatementId)){
            return null;
        }
        if (DATA_SCOPE_CACHE.containsKey(mappedStatementId)) {
            return DATA_SCOPE_CACHE.get(mappedStatementId);
        }
        String className = StringUtil.subBefore(mappedStatementId, StringPool.DOT, true);
        String methodName = StringUtil.subAfter(mappedStatementId, StringPool.DOT, true);
        if (Func.isAnyBlank(className, methodName)) {
            IGNORE_METHOD_CACHE.add(mappedStatementId);
            return null;
        }
        Class<?> aClass = ClassLoaderUtil.loadClass(Objects.requireNonNull(className));
        if (Func.isEmpty(aClass)) {
            IGNORE_METHOD_CACHE.add(mappedStatementId);
            return null;
        }
        List<Method> allMethods = Lists.newLinkedList();
        addAllMethod(Lists.newArrayList(aClass),allMethods);
        Method method = null;
        for (Method m : allMethods) {
            if(Func.equals(m.getName(),methodName)){
                method = m;
                break;
            }
        }
        if (Func.isEmpty(method)) {
            IGNORE_METHOD_CACHE.add(mappedStatementId);
            return null;
        }
        DataScope annotation = Func.getAnnotation(method, DataScope.class);
        if(annotation == null){
            IGNORE_METHOD_CACHE.add(mappedStatementId);
            return null;
        }
        DATA_SCOPE_CACHE.put(mappedStatementId, annotation);
        DATA_SCOPE_METHOD_CACHE.put(mappedStatementId, method);
        return annotation;
    }

}

创建通用处理器接口

1
2
3
4
5
6
7
8
9
10
11
public interface DataPermissionHandler {
    /**
     * 获取处理之后的where条件表达式，返回null时忽略
     * @param where 待执行的sql Where条件表达式，可能为null
     * @param mappedStatementId 全类名.方法名 一定不为null
     * @param dataScope 注解，进入处理方法时一定不为null
     * @param method 当前执行的查询方法，进入处理方法时一定不为null
     * @return 获取处理之后的where条件表达式
     */
    Expression getSqlSegment(Expression where, String mappedStatementId, DataScope dataScope, Method method);
}

创建默认的处理器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
public class DefaultDataPermissionHandler implements DataPermissionHandler{
    private static final String LOG_PREFIX = "默认数据权限处理器===";
    @Override
    public Expression getSqlSegment(Expression where, String mappedStatementId, DataScope dataScope, Method method) {
        // 两个字段都为空，无需验证
        if(Func.isNoneBlank(dataScope.createByColumnName(),dataScope.orgIdColumnName())){
            return null;
        }
        // 获取当前登陆人
        String currentUserId = Func.toStr(SecureUtil.getUserId());
        // 获取当前登陆人可以管理的机构ID
        Set<String> currentManageOrgIds = DataScopeUtil.getManageOrgIds();
        EqualsTo createByExpression = null;
        InExpression orgExpression = null;
        // 字段名不为空时处理
        if(Func.isNotBlank(dataScope.createByColumnName())){
            createByExpression = new EqualsTo();
            createByExpression.setLeftExpression(new Column(dataScope.createByColumnName()));
            createByExpression.setRightExpression(new StringValue(currentUserId));
        }
        // 字段名不为空时处理
        if(Func.isNotBlank(dataScope.orgIdColumnName())){
            orgExpression = new InExpression(new Column(dataScope.orgIdColumnName())
                    ,new ExpressionList(currentManageOrgIds.stream().map(LongValue::new).collect(Collectors.toList())));
        }
        // 两种都存在
        if(createByExpression != null && orgExpression != null){
            OrExpression dataPermission = new OrExpression(createByExpression,orgExpression);
            return Func.isNull(where)?dataPermission:new AndExpression(where,dataPermission);
        }
        // 只有创建人
        if(createByExpression != null){
            return Func.isNull(where)?createByExpression:new AndExpression(where,createByExpression);
        }
        // 只有机构
        if(orgExpression != null){
            return Func.isNull(where)?orgExpression:new AndExpression(where,orgExpression);
        }
        throw new RuntimeException(LOG_PREFIX+"系统逻辑处理异常，请检查");
    }
}

将拦截器注册到MybatisPlus

1
2
3
4
5
6
7
8
9
10
@Bean
@ConditionalOnMissingBean(MybatisPlusInterceptor.class)
public MybatisPlusInterceptor paginationInterceptor() {
    MybatisPlusInterceptor mybatisPlusInterceptor = new MybatisPlusInterceptor();
    // 添加数据权限拦截器(数据权限拦截器必须在分页之前)
    mybatisPlusInterceptor.addInnerInterceptor(new DataPermissionInterceptor());
    // 新版分页插件必须指定数据库方言 V3.5.1
    mybatisPlusInterceptor.addInnerInterceptor(new PaginationInnerInterceptor(DbType.MYSQL));
    return mybatisPlusInterceptor;
}

数据权限工具类

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
@Slf4j
public class DataScopeUtil {
    /**
    使用方法 ： orgService.list(DataScopeUtil.getDataScopeQuery(null,Org::getId))
     * 获取加入数据权限的query
     * @param createByColumn 创建人字段 , 传入null不做筛选
     * @param orgIdColumn 机构ID字段 ，传入null不做筛选
     * @param <T> 当前查询的实体类型
     * @return 获取包含数据权限的query
     */
    public static <T> LambdaQueryWrapper<T> getDataScopeQuery(SFunction<T,?> createByColumn,SFunction<T,?> orgIdColumn){
        if(SecureUtil.isAdministrator()){
            return Wrappers.lambdaQuery();
        }
        return Wrappers.<T>lambdaQuery()
                .eq(Func.notNull(createByColumn),createByColumn, SecureUtil.getUserId())
                .or()
                .in(Func.notNull(orgIdColumn),orgIdColumn, getManageOrgIds());
    }
}