MyBatis - 拦截器之修改预编译后的 SQL

因为 XML foreach 标签进行预编译后，对生成的 SQL 会有一些换行、空行、回车等不可见字符，于是就想能不能对生成后的 SQL 进行二次编辑，于是有了下文。

import java.lang.annotation.*;

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface SqlFormat {

}

public interface SqlDao extends DBaseMapper {

    
    @SqlFormat
    void createTable(SqlTablePO sqlTablePo);
}

import com.chinadaas.platform.dsp.pipeline.annotation.SqlFormat;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.type.TypeHandlerRegistry;
import java.lang.reflect.Method;
import java.text.DateFormat;
import java.util.*;
import java.util.regex.Matcher;


@Slf4j
@Intercepts({
        @Signature(type = org.apache.ibatis.executor.Executor.class, method = "update", args = {
                MappedStatement.class, Object.class
        })
})
public class SqlInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        log.info("SqlInterceptor intercept");

        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        SqlFormat sqlFormat = getSqlFormat(ms);
        BoundSql boundSql = ms.getBoundSql(args[1]);
        String oldSql = boundSql.getSql();
        log.info("OldSql: {}", oldSql);
        if (Objects.isNull(sqlFormat)) {
            return invocation.proceed();
        }

        // 生成新的 SQL
        String newSql = newSql(oldSql);

        // 重新塞回去
        BoundSql newBoundSql = new BoundSql(ms.getConfiguration(), newSql, boundSql.getParameterMappings(), boundSql.getParameterObject());
        MappedStatement newMs = newMappedStatement(ms, newBoundSql);
        for (ParameterMapping mapping : boundSql.getParameterMappings()) {
            String prop = mapping.getProperty();
            if (boundSql.hasAdditionalParameter(prop)) {
                newBoundSql.setAdditionalParameter(prop, boundSql.getAdditionalParameter(prop));
            }
        }
        args[0] = newMs;

        log.info("NewSql: {}", newBoundSql.getSql());

        return invocation.proceed();
    }

    private String newSql(String oldSql) {
        return oldSql.replace("n", "")
                .replace("r", "")
                .replace(" ", "")
                .replace("{Space}", " ");
    }

    private SqlFormat getSqlFormat(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();

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

    private MappedStatement newMappedStatement(MappedStatement ms, BoundSql newBoundSql) {
        MappedStatement.Builder builder = new
                MappedStatement.Builder(ms.getConfiguration(), ms.getId(), new BoundSqlSource(newBoundSql), 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();
    }

    @Override
    public Object plugin(Object target) {
        log.info("SqlInterceptor plugin");
        return Interceptor.super.plugin(target);
    }

    @Override
    public void setProperties(Properties properties) {
        log.info("SqlInterceptor setProperties");
        Interceptor.super.setProperties(properties);
    }

    class BoundSqlSource implements SqlSource {
        private BoundSql boundSql;

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

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

Ps：这边拦截只是针对 Update 操作，这个在上面自己配上 @Signature 即可，这个很简单就不再描述，对应填写的参数看下都会的！

@Bean
public SqlSessionFactory sqlSessionFactory() throws Exception {
    // 用 mybatis 的这里会有点区别, mybatis 用的是 SqlSessionFactoryBean
    MybatisSqlSessionFactoryBean sqlSessionFactoryBean = new MybatisSqlSessionFactoryBean();
    sqlSessionFactoryBean.setDataSource(dynamicDataSource());
    sqlSessionFactoryBean.setMapperLocations((new PathMatchingResourcePatternResolver()).getResources("classpath:mapper/*.xml"));
 
    // 注册 SQL 拦截器
    sqlSessionFactoryBean.setConfiguration(new MybatisConfiguration());
    sqlSessionFactoryBean.getConfiguration().addInterceptor(new SqlInterceptor());
 
    return sqlSessionFactoryBean.getObject();
}