mybatis拦截器做数据权限

package com.geely.dt.configuration;

import com.geely.dt.service.pipeline.*;
import com.geely.dt.service.pipeline.menuresouceperimissiondata.ResouceMenuData;
import com.geely.dt.service.pipeline.menuresouceperimissiondata.ResouceMenuDataPipeline;
import com.geely.dt.service.pipeline.menuresouceperimissiondata.ResouceMenuOperationData;
import com.geely.dt.service.pipeline.menuresouceperimissiondata.ResouceMenuUserData;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.*;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.stereotype.Component;
import org.apache.ibatis.plugin.*;

import java.util.List;
import java.util.stream.Collectors;

/**
 * Created by Jiyang.Zheng on 2022/4/18 14:32.
 */
@Slf4j
@Component
//mybatis 拦截顺序Executor -> StatementHandler->ParameterHandler->ResultSetHandler
//要在分页插件之前完成sql语句的修改 应拦截Executor
@Intercepts(
        {
                @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})
        }
)
public class DataPermissionInjectInterceptor implements Interceptor {


    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        //获取拦截下的mapper
        MappedStatement mappedStatement = (MappedStatement) invocation.getArgs()[0];
        //获取上一个拦截器传过来的sql
        BoundSql boundSql = mappedStatement.getBoundSql(invocation.getArgs()[1]);
        //根据指定权限点对原有sql进行修改
        String  mSql = modifyOrgSql(mappedStatement,boundSql,invocation);

        //通过反射修改sql语句
//            Field field = boundSql.getClass().getDeclaredField("sql");
//            field.setAccessible(true);
//            field.set(boundSql, mSql);

        return invocation.proceed();
    }

    /**
     * 根据权限点拼装对应sql
     * @return 拼装后的sql
     */
    private String modifyOrgSql(MappedStatement mappedStatement,BoundSql boundSql,Invocation invocation)  {
        String mSql = boundSql.getSql();
        if (!filterObject(mappedStatement.getId())) {
            return mSql;
        }
        String orgSQql = boundSql.getSql();

        Handler<Void, List<Integer>> handler =  new ResouceMenuData();
        Pipeline pipeline;
        if (StringUtils.equals("com.geely.dt.dao.ResourceMapper.selectDataPerimissonBusinessList",mappedStatement.getId())){
            pipeline = new ResouceMenuDataPipeline<>(handler)
                    .addHandler(new ResouceMenuOperationData())
                    .addHandler(new ResouceMenuUserData());
        }else {
            pipeline = new ResouceMenuDataPipeline<>(handler)
                    .addHandler(new ResouceMenuOperationData());
        }

        //是管理員
        if (!pipeline.before(null)){
            return orgSQql;
        }
        List<Integer> list = (List<Integer>) pipeline.execute(null);
        pipeline.after(list);
        log.info("原始SQL: {}", orgSQql);
        //没有任何数据
        if (CollectionUtils.isEmpty(list)){
            mSql = "select * from ( "+ orgSQql+" ) t where id in(-1)";
        }else {
            String val = list.stream().map(v -> v+"").collect(Collectors.joining(","));
            mSql = "select * from ( "+ orgSQql+" ) t where id in("+ val +")";
        }
        // 重新new一个查询语句对象
        BoundSql newBoundSql = new BoundSql(mappedStatement.getConfiguration(), mSql,
                boundSql.getParameterMappings(), boundSql.getParameterObject());
        // 把新的查询放到statement里
        MappedStatement newMs =  newMappedStatement (mappedStatement, new BoundSqlSqlSource(newBoundSql));
        for (ParameterMapping mapping : boundSql.getParameterMappings()) {
            String prop = mapping.getProperty();
            if (boundSql.hasAdditionalParameter(prop)) {
                newBoundSql.setAdditionalParameter(prop, boundSql.getAdditionalParameter(prop));
            }
        }
        Object[] queryArgs = invocation.getArgs();
        queryArgs[0] = newMs;
        log.info("改写的SQL: {}", mSql);
        return mSql;
    }

    /**
     * 定义一个内部辅助类，作用是包装 SQL
     */
    class BoundSqlSqlSource implements SqlSource {
        private BoundSql boundSql;
        public BoundSqlSqlSource(BoundSql boundSql) {
            this.boundSql = boundSql;
        }
        public BoundSql getBoundSql(Object parameterObject) {
            return boundSql;
        }

    }

    private MappedStatement newMappedStatement (MappedStatement ms, SqlSource newSqlSource) {
        MappedStatement.Builder builder = new
                MappedStatement.Builder(ms.getConfiguration(), ms.getId(), newSqlSource, ms.getSqlCommandType());
        builder.resource(ms.getResource());
        builder.fetchSize(ms.getFetchSize());
        builder.statementType(ms.getStatementType());
        builder.keyGenerator(ms.getKeyGenerator());
        if (ms.getKeyProperties() != null && ms.getKeyProperties().length > 0) {
            builder.keyProperty(ms.getKeyProperties()[0]);
        }
        builder.timeout(ms.getTimeout());
        builder.parameterMap(ms.getParameterMap());
        builder.resultMaps(ms.getResultMaps());
        builder.resultSetType(ms.getResultSetType());
        builder.cache(ms.getCache());
        builder.flushCacheRequired(ms.isFlushCacheRequired());
        builder.useCache(ms.isUseCache());
        return builder.build();
    }


    private static List<String> filterIds = Lists.newArrayList();
    static {
        filterIds.add("com.geely.dt.dao.ResourceMapper.selectDataPerimissonBackList");
        filterIds.add("com.geely.dt.dao.ResourceMapper.selectDataPerimissonBusinessList");
    }


    /**
     * 默认拦截所有SQL
     * @param
     * @return boolean
     */
    public boolean filterObject(String id) {
        if (filterIds.contains(id)){
            return true;
        }
        return false;
    }
}