MyBatis拦截器 Interceptor 实现多数据源切换

实现思路

        对于无事务的数据库 DQL操作，检查其 sqlCommandType、是否使用了 for update、或者检查是否存在指定注解。

前置知识

        AbstractRoutingDataSource

http://t.csdn.cn/n68tvhttp://t.csdn.cn/n68tv数据源配置与切换

public enum DataSources {
    // 主库-读数据源
    master,
    // 主库-写数据源
    slave,
    // 多数据源
    readmore
}

package com.gateway.admin.datasources;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import javax.sql.DataSource;
import java.util.Map;

/**
 * 数据源路由器
 */
public class DataSourceRouter extends AbstractRoutingDataSource {
    // 也可以指定 ThreadLocal 的 initialValue 的具体实现
    private static final ThreadLocal<String> contextHolder = new ThreadLocal<>();

    public DynamicDataSource(DataSource defaultTargetDataSource, Map<Object, Object> targetDataSources) {
        super.setDefaultTargetDataSource(defaultTargetDataSource);
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }
    /**
     *将targetDataSources映射中指定的给定查找键对象解析为用于与当前查找键匹配的实际查找键。
     */
    @Override
    protected Object determineCurrentLookupKey() {
        return contextHolder.get();
    }
    // 设置数据源,参数类型是上面生命的枚举
    public static void setDataSource(DataSources dataSource) {
        contextHolder.set(dataSource);
    }

    // 在 finally 中及时 remove    
    public static void clearDataSource() {
        contextHolder.remove();
    }

}

package com.gateway.admin.datasources;

import com.alibaba.druid.spring.boot.autoconfigure.DruidDataSourceBuilder;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * 多数据源配置类
 */
@Configuration
public class DynamicDataSourceConfig {

	//如果ioc容器中,同一个类型有多个bean,则bean的名称为方法的名称
    @Bean
    @ConfigurationProperties("spring.datasource.druid.first")
    public DataSource firstDataSource() {
        return DruidDataSourceBuilder.create().build();
    }

    @Bean
    @ConfigurationProperties("spring.datasource.druid.second")
    public DataSource secondDataSource() {
        return DruidDataSourceBuilder.create().build();
    }

    @Bean
    @ConfigurationProperties("spring.datasource.druid.three")
    public DataSource threeDataSource() {
        return DruidDataSourceBuilder.create().build();
    }

    @Bean
    @ConfigurationProperties("spring.datasource.druid.four")
    public DataSource fourDataSource() {
        return DruidDataSourceBuilder.create().build();
    }

    @Bean("dataSource")
    @Primary
    public DataSourceRouter dataSource(DataSource firstDataSource, DataSource secondDataSource, DataSource threeDataSource, DataSource fourDataSource) {
        Map<Object, Object> targetDataSources = new HashMap<>();
        targetDataSources.put(DataSourceNames.master, master);
        targetDataSources.put(DataSourceNames.slave, slave);
        targetDataSources.put(DataSourceNames.readmore, readmore);
        
        // DataSourceNames.master, master 设置为默认数据源头
        return new DataSourceRouter(firstDataSource, targetDataSources);
    }
}

yml 配置

spring:
    datasource:
        type: com.alibaba.druid.pool.DruidDataSource
        driverClassName: com.mysql.jdbc.Driver
        druid:
            first:  #db1
                url: jdbc:mysql://127.0.0.1:3306/db1?allowMultiQueries=true&useUnicode=true&characterEncoding=UTF-8&useSSL=false
                username: root
                password: root
            second:  #db2
                url: jdbc:mysql://127.0.0.1:3306/db2?allowMultiQueries=true&useUnicode=true&characterEncoding=UTF-8&useSSL=false
                username: root
                password: root
            three:  #db3
                url: jdbc:mysql://127.0.0.1:3306/db3?allowMultiQueries=true&useUnicode=true&characterEncoding=UTF-8&useSSL=false
                username: root
                password: root

具体实现

annotation

@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface DataSource {


    DataSources value() default DataSources.master;
}

 DynamicDataSourceInterceptor

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


    /**
     * 可以识别为 添加、更新、删除类型的 SQL 语句
     */
    public static final List<SqlCommandType> UPDATE_SQL_LIST = Arrays.asList(SqlCommandType.INSERT, SqlCommandType.UPDATE, SqlCommandType.DELETE);

    /**
     * SQL 语句中出现的悲观锁标识
     */
    private static final String LOCK_KEYWORD = "for update";

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        // 通过 invocation 获取 MappedStatement 与 拦截方法的形参信息
        Object[] objects = invocation.getArgs();
        MappedStatement ms = (MappedStatement) objects[0];

        // 通过反射检查要执行的方法，如果标注了 @DataSource 则检查其 value
        String clazzStr = ms.getId().substring(0, ms.getId().lastIndexOf("."));
        String methodStr = ms.getId().substring(ms.getId().lastIndexOf(".") + 1);
        // 由于 mybatis 同一个接口方法不能重载
        Method[] mapperMethods = Class.forName(clazzStr).getMethods();
        Method targetMethod = null;
        for (Method mapperMethod : mapperMethods) {
            if (mapperMethod.getName().equals(methodStr)) {
                targetMethod = mapperMethod;
                break;
            }
        }
        DataSources dataSourceAnnotationValue = null;
        if (targetMethod != null && targetMethod.getAnnotation(DataSource.class) != null) {
            dataSourceAnnotationValue = targetMethod.getAnnotation(DataSource.class).value();
        }

        // 获取 sqlCommandType
        SqlCommandType sqlCommandType = ms.getSqlCommandType();

        // 获取 SQL
        BoundSql boundSql = ms.getSqlSource().getBoundSql(objects[1]);
        String sql = boundSql.getSql().toLowerCase(Locale.CHINA).replace("[\\t\\n\\r]", " ");

        if (dataSourceAnnotationValue == DataSources.read && sqlCommandType.equals(SqlCommandType.SELECT)) {

            DataSourceTypeManager.set(DataSources.slave);

        } else if (dataSourceAnnotationValue == DataSources.write ||
                UPDATE_SQL_LIST.contains(sqlCommandType) ||
                sql.contains(LOCK_KEYWORD)) {

            DataSourceTypeManager.set(DataSources.slave);

        } else {
            DataSourceTypeManager.set(DataSources.readmore);

        }

        Object proceed;
        try {
            proceed = invocation.proceed();
        } catch (Throwable t) {
            throw t;
        } finally {
            DataSourceTypeManager.reset();
        }
        return proceed;
    }

    @Override
    public Object plugin(Object target) {

        if (target instanceof Executor) {

            return Plugin.wrap(target, this);
        } else {
            return target;
        }
    }

    @Override
    public void setProperties(Properties properties) {

    }
}

mybatis-config.xml

.....
    <plugins>
        <plugin interceptor="interceptor.DynamicDataSourceInterceptor" />
    </plugins>

上盘硬菜，@Transaction源码深度解析 | Spring系列第48篇 (qq.com)