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

Posted 2022-07-02 放羊的牧码

故事背景

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

注解

import java.lang.annotation.*;

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

Dao 层

/**
 * @author Lux Sun
 * @date 2022/3/17
 */
public interface SqlDao extends DBaseMapper<SqlFieldPO> 

    /**
     * 创建表
     * @param sqlTablePo
     */
    @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;

/**
 * @author Lux Sun
 * @date 2022/6/27
 */
@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();