MyBatis自定义SQL拦截器示例详解

// 全局开启
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.TYPE)
@Documented
@Import(MyBatisSqlInterceptorConfiguration.class)
public @interface EnableSqlInterceptor {

}

// 自定义注解
@Target({ElementType.METHOD })
@Retention(RetentionPolicy.RUNTIME)
public @interface DataScope {

}

public class MyBatisSqlInterceptorConfiguration implements ApplicationContextAware {

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        SqlSessionFactory sqlSessionFactory = applicationContext.getBean(SqlSessionFactory.class);
        sqlSessionFactory.getConfiguration().addInterceptor(new MyBatisInterceptor());
    }
}

@Slf4j
@Intercepts(
        {
                @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
                @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}),
        })
public class MyBatisInterceptor implements Interceptor {

    private static final Logger LOGGER = LoggerFactory.getLogger(MyBatisInterceptor.class);

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        // TODO Auto-generated method stub
        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        Object parameter = args[1];
        RowBounds rowBounds = (RowBounds) args[2];
        ResultHandler resultHandler = (ResultHandler) args[3];
        Executor executor = (Executor) invocation.getTarget();
        CacheKey cacheKey;
        BoundSql boundSql;
        //由于逻辑关系，只会进入一次
        if (args.length == 4) {
            //4 个参数时
            boundSql = ms.getBoundSql(parameter);
            cacheKey = executor.createCacheKey(ms, parameter, rowBounds, boundSql);
        } else {
            //6 个参数时
            cacheKey = (CacheKey) args[4];
            boundSql = (BoundSql) args[5];
        }
        DataScope dataScope = getDataScope(ms);
        if (Objects.nonNull(dataScope)) {
            String origSql = boundSql.getSql();
            log.info("origSql : {}", origSql);
            // 组装新的 sql
            // todo you weaving business
            String newSql = origSql + " limit 1";

            // 重新new一个查询语句对象
            BoundSql newBoundSql = new BoundSql(ms.getConfiguration(), newSql,
                    boundSql.getParameterMappings(), boundSql.getParameterObject());

            // 把新的查询放到statement里
            MappedStatement newMs = newMappedStatement(ms, new BoundSqlSource(newBoundSql));
            for (ParameterMapping mapping : boundSql.getParameterMappings()) {
                String prop = mapping.getProperty();
                if (boundSql.hasAdditionalParameter(prop)) {
                    newBoundSql.setAdditionalParameter(prop, boundSql.getAdditionalParameter(prop));
                }
            }

            args[0] = newMs;
            if (args.length == 6) {
                args[5] = newMs.getBoundSql(parameter);
            }
        }
        LOGGER.info("mybatis intercept sql:{},Mapper方法是：{}", boundSql.getSql(), ms.getId());


        return invocation.proceed();
    }

    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 DataScope getDataScope(MappedStatement mappedStatement) {
        String id = mappedStatement.getId();
        // 获取 Class Method
        String clazzName = id.substring(0, id.lastIndexOf('.'));
        String mapperMethod = id.substring(id.lastIndexOf('.') + 1);

        Class<?> clazz;
        try {
            clazz = Class.forName(clazzName);
        } catch (ClassNotFoundException e) {
            return null;
        }
        Method[] methods = clazz.getMethods();

        DataScope dataScope = null;
        for (Method method : methods) {
            if (method.getName().equals(mapperMethod)) {
                dataScope = method.getAnnotation(DataScope.class);
                break;
            }
        }
        return dataScope;
    }

    @Override
    public Object plugin(Object target) {
        // TODO Auto-generated method stub
        LOGGER.info("MysqlInterCeptor plugin>>>>>>>{}", target);
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
        // TODO Auto-generated method stub
        String dialect = properties.getProperty("dialect");
        LOGGER.info("mybatis intercept dialect:>>>>>>>{}", dialect);
    }

    /**
     * 定义一个内部辅助类，作用是包装 SQL
     */
    class BoundSqlSource implements SqlSource {
        private BoundSql boundSql;

        public BoundSqlSource(BoundSql boundSql) {
            this.boundSql = boundSql;
        }

        public BoundSql getBoundSql(Object parameterObject) {
            return boundSql;
        }

    }

}

@Mapper
public interface OrderMapper {

   @Select("select 1 ")
   @DataScope
   Intger selectOne();

}