java-mybatis-自定义interceptor拦截器-获取原生sql补全公共参数

java-mybatis-自定义interceptor拦截器-获取原生sql补全公共参数

环境 pgsql + jdk1.8 + mybatis + springboot

一、引言

数据使用PGSQL存储，由于业务场景扩充，使用固定的Schema不能满足业务存储需求，亟需修改代码逻辑将原来在代码中写死的Schema名称改为动态可变形式。
最简单的方法是将所有引用 mySchema 的地方使用 ${schameName} 替换，再修改调用方法调用处，将schemaName传入。

--原查询脚本
select * from mySchema.myTable;
--替换后的形式
select * from ${mySchema}.myTable;

之所以用${mySchema}占位符，是因为拼接表名不能使用参数化方式传参；
好吧，全局搜下需要修改的位置，看看有多少：

--查询结果为：
Find in Files 300+ matches in 30+ files

分析了匹配的文件，发现引用位置有的在 xml 中，有的在 @Select、@Update等注解中，果断放弃！！！
想到了用mybatis的拦截器，不需要修改调用方法的形参，使用非侵入性的修改就可以完成参数注入；
先看下拦截器相关基础资料。

二、四种mybatis拦截器

Mybatis提供了四个可拦截对象是：

• Executor sql的内部执行器
• ParameterHandler 拦截参数的处理
• StatementHandler 拦截sql的构建
• ResultSetHandler 拦截结果的处理

这4中不同类型的拦截器的拦截顺序为从上到下的顺序为：
Executor -> ParameterHandler -> StatementHandler -> ResultSetHandler
如果相同类型的拦截器,比如Executor类型的拦截器有两个,则执行顺序为将拦截器添加到SqlSessionFactory的逆向顺序执行；
比如SqlSessionFactory中先添加了Executor类型的A拦截器,在添加了Executor类型的B拦截器,则会先执行B拦截器,再执行A拦截器；

声明拦截器

• 使用@Intercepts、@Signature注解声明拦截器

@Component
@Intercepts({
    @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
})
public class SchemaParamsterInterceptor implements Interceptor {
}

• @Signature注解参数说明

• type 四个可拦截类的class对象；
• method 类中的方法名称，入上述的prepare就是StatementHandler中的prepare方法；
• args 方法的形参类型，数量类型与拦截类中的方法声明一致；

注册拦截器

• 使用 @Configuration配置注解

@Configuration
public class MybatisConfig {
    @Autowired
    private List<SqlSessionFactory> sqlSessionFactoryList;
    /**
     * mybatis 拦截器注册
     */
    @PostConstruct
    public void addSqlInterceptor() {
        SchemaParamsterInterceptor interceptor = new SchemaParamsterInterceptor();
        for (SqlSessionFactory sqlSessionFactory : sqlSessionFactoryList) {
            sqlSessionFactory.getConfiguration().addInterceptor(interceptor);
        }
    }
}

三、动态参数注入方案：

方案一：使用正常的sql占位符${…}

思路是在sql定义中使用标准占位符 ${schemaName}，在拦截器中判断原始SQL中是否出现${schemaName}字符，出现了就注入schemaName参数。
优点：可以兼容标准调用，开发者对于schameName可传可不传，不传就可以自动注入，显示传入则使用传入的参数；

• 使用ParameterHandler拦截器遇到的问题：

• sql定义中写了${schemaName}占位符但是没传入，在Executor拦截器中就异常了，根本走不到ParameterHandler拦截器；
• 那试着用Executor拦截器；

• 使用Executor拦截器遇到的问题：

• 测试失败，因为传入参数类型的类型不固定（map、pojo、基础类型）不方便判断和注入；

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

    @Override
    public Object plugin(Object target) {
        return Interceptor.super.plugin(target);
    }

    @Override
    public void setProperties(Properties properties) {
        Interceptor.super.setProperties(properties);
    }
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        // 拦截 Executor 的 query 方法 生成sql前将 任意参数 设置到实体中
        if (invocation.getTarget() instanceof Executor ) {
            //&& "query".equals(invocation.getMethod().getName())
            return invokeQuery(invocation);
        }
        return null;
    }

    /** 获取原始sql */
    private String getRawSQL(Invocation invocation) throws NoSuchFieldException, IllegalAccessException {
        //反射获取 SqlSource 对象，通过此对象获取原始SQL
        MappedStatement ms = (MappedStatement) invocation.getArgs()[0];
        /** SqlSource{@link org.apache.ibatis.mapping.SqlSource}的实现类比较多，不方便在所有实现类中解析原始SQL*/
        SqlSource sqlSource = ms.getSqlSource();
        //通过MappedStatement.SqlSource对象获取原生sql不太靠谱！！！
        if(sqlSource instanceof DynamicSqlSource) {
            DynamicSqlSource dynamicSqlSource = (DynamicSqlSource) ms.getSqlSource();
            if (dynamicSqlSource == null)
                return null;
            //反射获取 TextSqlNode 对象
            Field sqlNodeField = dynamicSqlSource.getClass().getDeclaredField("rootSqlNode");
            sqlNodeField.setAccessible(true);
            TextSqlNode rootSqlNode = (TextSqlNode) sqlNodeField.get(dynamicSqlSource);
            //反射获取原生sql
            Field textField = rootSqlNode.getClass().getDeclaredField("text");
            textField.setAccessible(true);
            String sql = String.valueOf(textField.get(rootSqlNode));
            return sql;
        }
        if(sqlSource instanceof RawSqlSource) {
            RawSqlSource rawSqlSource = (RawSqlSource) ms.getSqlSource();
            if (rawSqlSource == null)
                return null;
            //反射获取 TextSqlNode 对象
            Field sqlSourceField = rawSqlSource.getClass().getDeclaredField("sqlSource");
            sqlSourceField.setAccessible(true);
            StaticSqlSource staticSqlSource = (StaticSqlSource) sqlSourceField.get(rawSqlSource);
            //反射获取原生sql
            Field sqlField = staticSqlSource.getClass().getDeclaredField("sql");
            sqlField.setAccessible(true);
            String sql = String.valueOf(sqlField.get(staticSqlSource));
            return sql;
        }
        return null;
    }

    private Object invokeQuery(Invocation invocation) throws Exception {
        //todo 按需添加注入參數提高性能
//        String sql = getRawSQL(invocation);
//        if(StringUtils.isBlank(sql) || sql.indexOf(schemaParamsPlaceholder)==-1)
//            return null;

        Executor executor = (Executor) invocation.getTarget();

        // 获取第一个参数
        MappedStatement ms = (MappedStatement) invocation.getArgs()[0];
        // mybatis的参数对象
        Object paramObj = invocation.getArgs()[1];
        if (paramObj == null) {
            MapperMethod.ParamMap<Object> param = new MapperMethod.ParamMap<>();
            paramObj = param;
        }
        //执行脚本
        processParam(paramObj);
        RowBounds rowBounds = (RowBounds)invocation.getArgs()[2];
        ResultHandler resultHandler  = (ResultHandler)invocation.getArgs()[3];
        return executor.query(ms, paramObj,rowBounds,resultHandler);
    }
    /** 处理参数对象 */
    private void processParam(Object parameterObject) throws IllegalAccessException, InvocationTargetException {
        //如果是map且map的key中没有需要的参数，则添加到参数map中
        if (parameterObject instanceof Map) {
            String schemaName = "mySchema";
            ((Map) parameterObject).putIfAbsent("schemaName", schemaName);
            return;
        }
    }
}

方案二：使用非标准占位符

思路：使用自定的占位符，这样就可以跳过参数定义校验，在sql预处理时将占位符替换为正确的schema名；
优点：可以兼容xml和通过注解定义的sql,支持所有类型sql如：select、insert、delete、update；
缺点：不能兼容显示传入的参数，所有执行的sql都会被拦截；

/***
 * schema参数动态注入
 */
@Log4j
@Component
@Intercepts({
        @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
})
public class SchemaParamsterInterceptor implements Interceptor {
    /** SQL中的占位符 */
    private static final String schemaPlaceholder = "_schemaName";
    

    @Override
    public Object plugin(Object target) {
        return Interceptor.super.plugin(target);
    }

    @Override
    public void setProperties(Properties properties) {
        Interceptor.super.setProperties(properties);
    }
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        //
        if(invocation.getTarget() instanceof  StatementHandler){
            if("prepare".equals(invocation.getMethod().getName()))
                return invokeStatementHandlerPrepare(invocation);
        }
        return null;
    }

    private Object invokeStatementHandlerPrepare(Invocation invocation) throws ClassNotFoundException, NoSuchFieldException, IllegalAccessException, InvocationTargetException {
        StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
        BoundSql boundSql = statementHandler.getBoundSql();
        String sql = boundSql.getSql();
        log.debug("prepare~~~~~~~~~~~~~~~begin");
        System.out.println(sql);
        if(StringUtils.isNotEmpty(sql) && sql.indexOf(schemaPlaceholder)>-1){
            String adminSchema = "mySchema";
            sql = sql.replaceAll(schemaPlaceholder,adminSchema);
            //通过反射回写
            Field sqlNodeField = boundSql.getClass().getDeclaredField("sql");
            sqlNodeField.setAccessible(true);
            sqlNodeField.set(boundSql,sql);
            log.debug("prepare~~~~~~~~~~~~~~~replace");
            log.debug(sql);
        }
        log.debug("prepare~~~~~~~~~~~~~~~end");


        return invocation.proceed();
    }
}

四、结论

使用方案二，完美解决了动态替换schame的需求。
项目中使用 Replace in Files 将原来的固定schame名称，替换为 _schemaName 占位符。
代码完美运行！！！