MyBatis-Plus数据权限插件

MyBatis-Plus数据权限插件

一、实现思路和注意事项

• 实现思路

AOP注解和线程上下文变量传递参数到Mybatis-Plus SQL拦截器，使用JSqlParser AST SQL语法解析树操作Where动态添加过滤条件

• 注意事项

@DataScope注解一般用于Service层或者DAO层(Mapper)，用于Service层时，若方法中存在多个查询语句均会拼接数据权限过滤条件。

二、源码

1.pom.xml

    <dependencies>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>
        <dependency>
            <groupId>com.baomidou</groupId>
            <artifactId>mybatis-plus</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-configuration-processor</artifactId>
            <optional>true</optional>
        </dependency>
    </dependencies>

2.spring.factories

org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
com.example.datascope.autoconfigure.DataScopeAutoConfiguration

3.基础代码

• DataScope

package com.example.datascope.annotation;

import com.example.datascope.enums.DataScopeType;

import java.lang.annotation.*;

/**
 * 数据权限过滤注解
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface DataScope {

    /**
     * 范围权限类别
     */
    public DataScopeType type() default DataScopeType.DATA_SCOPE_SELF;

    /**
     * 部门表的别名
     */
    public String deptAlias() default "d";

    /**
     * 用户表的别名
     */
    public String userAlias() default "u";
}

• EnableDataScopeAspect

package com.example.datascope.annotation;

import com.example.datascope.aspectj.DataScopeAspect;
import org.springframework.context.annotation.Import;

import java.lang.annotation.*;

/**
 * 启用数据权限AOP
 */
@Target(ElementType.TYPE)
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Import(DataScopeAspect.class)
public @interface EnableDataScopeAspect {

}

• DataScopeAspect

package com.example.datascope.aspectj;

import com.example.datascope.annotation.DataScope;
import com.example.datascope.context.DataScopeContextHolder;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.stereotype.Component;

/**
 * 数据过滤处理
 */
@Aspect
@Component
public class DataScopeAspect {

    @Before("@annotation(datascope)")
    public void doBefore(JoinPoint point, DataScope datascope) throws Throwable {
        resetContextHolders();
        initContextHolders(datascope);
    }

    protected void initContextHolders(DataScope dataScope) {
        DataScopeContextHolder.set(dataScope);
    }

    private void resetContextHolders() {
        DataScopeContextHolder.remove();
    }

}

• DataScopeAutoConfiguration

package com.example.datascope.autoconfigure;

import com.example.datascope.annotation.EnableDataScopeAspect;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Configuration;

@EnableDataScopeAspect
@Configuration(proxyBeanMethods = false)
@ConditionalOnProperty(prefix = "datascope", value = "enable", havingValue = "true", matchIfMissing = true)
public class DataScopeAutoConfiguration {

}

• DataScopeContextHolder

package com.example.datascope.context;

import com.example.datascope.annotation.DataScope;

/**
 * DataScope上下文对象
 */
public abstract class DataScopeContextHolder {

    private static final ThreadLocal<DataScope> context = new InheritableThreadLocal<>();

    public static void set(DataScope dataScope) {
        context.set(dataScope);
    }

    public static DataScope get() {
        return context.get();
    }

    public static void remove() {
        context.remove();
    }

}

• DataScopeType

package com.example.datascope.enums;

import java.util.Objects;

public enum DataScopeType {
    /**
     * 全部数据权限
     */
    DATA_SCOPE_ALL("1"),
    /**
     * 自定数据权限
     */
    DATA_SCOPE_CUSTOM("2"),
    /**
     * 部门数据权限
     */
    DATA_SCOPE_DEPT("3"),
    /**
     * 部门及以下数据权限
     */
    DATA_SCOPE_DEPT_AND_CHILD("4"),
    /**
     * 仅本人数据权限
     */
    DATA_SCOPE_SELF("5");

    private String code;

    public String getCode() {
        return code;
    }

    DataScopeType(String code) {
        this.code = code;
    }

    public static DataScopeType of(String code) {

        Objects.requireNonNull(code, "数据范围权限类型不允许为空");

        for (DataScopeType dataScopeType : DataScopeType.values()) {
            if (dataScopeType.getCode().equals(code)) {
                return dataScopeType;
            }
        }

        throw new IllegalArgumentException(String.format("未识别的数据范围权限类型值[%s]", code));
    }


}

4.DataScopeInnerInterceptor

/*
 * Copyright (c) 2011-2022, baomidou (jobob@qq.com).
 *
 * 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
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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.example.datascope.plugins;

import com.example.datascope.enums.DataScopeType;
import com.example.datascope.annotation.DataScope;
import com.example.datascope.context.DataScopeContextHolder;
import com.example.datascope.handler.DataScopeHandler;
import com.baomidou.mybatisplus.core.plugins.InterceptorIgnoreHelper;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.baomidou.mybatisplus.core.toolkit.ExceptionUtils;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import com.baomidou.mybatisplus.extension.plugins.inner.TenantLineInnerInterceptor;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.*;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.logging.Log;
import org.apache.ibatis.logging.LogFactory;
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;

/**
 * Mybatis-Plus数据权限插件
 *
 * 
 * @see TenantLineInnerInterceptor
 */
public class DataScopeInnerInterceptor implements InnerInterceptor {

    private static final Log logger = LogFactory.getLog(DataScopeInnerInterceptor.class);

    private DataScopeHandler dataScopeHandler;

    public DataScopeInnerInterceptor(DataScopeHandler dataScopeHandler) {
        this.dataScopeHandler = dataScopeHandler;
    }

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

        // 原始SQL
        String originalSql = mpBs.sql();
        if (logger.isDebugEnabled()) {
            logger.debug("Original SQL: " + originalSql);
        }

        try {
            Statement statement = CCJSqlParserUtil.parse(originalSql);

            if (statement instanceof Select) {
                Select select = (Select) statement;
                // 解析SQL
                this.processSelect(select);
                final String parserSQL = statement.toString();
                mpBs.sql(parserSQL);
                if (logger.isDebugEnabled()) {
                    logger.debug("parser sql: " + parserSQL);
                }
            }
        } catch (JSQLParserException e) {
            throw ExceptionUtils.mpe("Failed to process, Error SQL: %s", e, originalSql);
        }
    }

    protected void processSelect(Select select) {
        processSelectBody(select.getSelectBody());
        List<WithItem> withItemsList = select.getWithItemsList();
        if (!CollectionUtils.isEmpty(withItemsList)) {
            withItemsList.forEach(this::processSelectBody);
        }
    }

    protected void processSelectBody(SelectBody selectBody) {
        if (selectBody == null) {
            return;
        }
        if (selectBody instanceof PlainSelect) {
            processPlainSelect((PlainSelect) selectBody);
        } else if (selectBody instanceof WithItem) {
            // With关键字
            WithItem withItem = (WithItem) selectBody;
            /**
             * jsqlparser 4.3版本 使用 {@code withItem.getSubSelect().getSelectBody())} 代替 {@code withItem.getSelectBody()}
             */
            processSelectBody(withItem.getSelectBody());
        } else {
            // 集合操作 UNION(并集) MINUS(差集)
            SetOperationList operationList = (SetOperationList) selectBody;
            List<SelectBody> selectBodyList = operationList.getSelects();
            if (CollectionUtils.isNotEmpty(selectBodyList)) {
                selectBodyList.forEach(this::processSelectBody);
            }
        }
    }

    /**
     * 处理 PlainSelect
     */
    protected void processPlainSelect(PlainSelect plainSelect) {
        DataScope dataScope = DataScopeContextHolder.get();
        if (dataScope == null) {
            return;
        } else {
            DataScopeType type = dataScope.type();
            try {
                switch (type) {
                    case DATA_SCOPE_ALL:
                        dataScopeHandler.setWhereForAll(plainSelect, dataScope);
                        break;
                    case DATA_SCOPE_CUSTOM:
                        dataScopeHandler.setWhereForCustom(plainSelect, dataScope);
                        break;
                    case DATA_SCOPE_DEPT:
                        dataScopeHandler.setWhereForDept(plainSelect, dataScope);
                        break;
                    case DATA_SCOPE_DEPT_AND_CHILD:
                        dataScopeHandler.setWhereForDeptAndChild(plainSelect, dataScope);
                        break;
                    case DATA_SCOPE_SELF:
                        dataScopeHandler.setWhereForSelf(plainSelect, dataScope);
                        break;
                }
            } catch (JSQLParserException e) {
                throw ExceptionUtils.mpe("Failed to process, Error SQL: %s", e);
            }
        }

    }

}

5.DataScopeHandler

package com.example.datascope.handler;

import com.example.datascope.annotation.DataScope;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.statement.select.PlainSelect;

/**
 * 数据权限处理逻辑
 * 参考: https://gitee.com/baomidou/mybatis-mate-examples/blob/master/mybatis-mate-datascope/src/main/java/mybatis/mate/datascope/config/DataScopeConfig.java
 */
public interface DataScopeHandler {
    /**
     * 全部数据权限
     *
     * @param plainSelect
     * @param dataScope   数据范围注解
     * @throws JSQLParserException SQL解析异常
     */
    default void setWhereForAll(PlainSelect plainSelect, DataScope dataScope) throws JSQLParserException {
        // do nothing
    }

    /**
     * 自定数据权限
     *
     * @param plainSelect
     * @param dataScope   数据范围注解
     * @throws JSQLParserException SQL解析异常
     */
    default void setWhereForCustom(PlainSelect plainSelect, DataScope dataScope) throws JSQLParserException {
        throw new UnsupportedOperationException("暂不支持的数据权限类型");
    }

    /**
     * 部门数据权限
     *
     * @param plainSelect
     * @param dataScope   数据范围注解
     * @throws JSQLParserException SQL解析异常
     */
    void setWhereForDept(PlainSelect plainSelect, DataScope dataScope) throws JSQLParserException;

    /**
     * 部门及以下数据权限
     *
     * @param plainSelect
     * @param dataScope   数据范围注解
     * @throws JSQLParserException SQL解析异常
     */
    void setWhereForDeptAndChild(PlainSelect plainSelect, DataScope dataScope) throws JSQLParserException;

    /**
     * 仅本人数据权限
     *
     * @param plainSelect
     * @param dataScope   数据范围注解
     * @throws JSQLParserException SQL解析异常
     */
    void setWhereForSelf(PlainSelect plainSelect, DataScope dataScope) throws JSQLParserException;

}

三、插件集成

• DataScopeHandler接口

根据实际情况，结合用户表和部门表结构，自实现DataScopeHandler接口，demo如下
可参考 mybatis-mate-datascope

@Configuration
public class DataScopeHandlerAutoConfigure {

    @Bean
    DataScopeHandler dataScopeHandler() {
        return new DataScopeHandler() {

            // TODO 从上下文中获取用户ID和部门ID

            @Override
            public void setWhereForDept(PlainSelect plainSelect, DataScope dataScope) {
                String column = String.format("%s.dept_id", dataScope.deptAlias());
                EqualsTo expression = new EqualsTo();
                expression.setLeftExpression(new Column(column));
                expression.setRightExpression(new StringValue("1"));
                setWhere(plainSelect, expression);
            }

            @Override
            public void setWhereForDeptAndChild(PlainSelect plainSelect, DataScope dataScope) throws JSQLParserException {
                String sqlPart = String.format("%s.dept_id IN ( SELECT dept_id FROM sys_dept WHERE dept_id = '%s' or find_in_set( '%s' , ancestors ) )", dataScope.deptAlias(), "1", "1");
                Expression expression = CCJSqlParserUtil.parseCondExpression(sqlPart);
                setWhere(plainSelect, expression);
            }

            @Override
            public void setWhereForSelf(PlainSelect plainSelect, DataScope dataScope) {
                String column = String.format("%s.user_id", dataScope.userAlias());
                EqualsTo expression = new EqualsTo();
                expression.setLeftExpression(new Column(column));
                expression.setRightExpression(new LongValue(1L));
                setWhere(plainSelect, expression);
            }

            private void setWhere(PlainSelect plainSelect, Expression expression) {
                Expression where = plainSelect.getWhere();
                if (where == null) {
                    // 不存在 where 条件
                    plainSelect.setWhere(new Parenthesis(expression));
                } else {
                    // 存在 where 条件 and 处理
                    plainSelect.setWhere(new AndExpression(plainSelect.getWhere(), expression));
                }
            }

        };
    }

}

• 启用DataScopeInnerInterceptor插件

@Configuration
public class MybatisPlusAutoConfigure {
    @Bean
    public MybatisPlusInterceptor mybatisPlusInterceptor() {
        MybatisPlusInterceptor interceptor = new MybatisPlusInterceptor();
        interceptor.addInnerInterceptor(new DataScopeInnerInterceptor(dataScopeHandler));
        return interceptor;
    }
}

• 启用@DataScope AOP注解扫描(默认启用)

datascope
  enable: true

四、使用示例

@DataScope(type = DataScopeType.DATA_SCOPE_DEPT_AND_CHILD, userAlias = "u", deptAlias = "d")

五、参考

TenantLineInnerInterceptor

mybatis-mate-datascope

RuoYi-Vue

Mybatis-Plus入门系列(18) -基于注解的动态数据权限实现方案