Mybatis数据源切换

记录一下自己通过 AbstractRoutingDataSource 和 Mybatis 插件实现读写分离的代码（参考了网上的一些思路）

一、AbstractRoutingDataSource

Spring 中提供了 AbstractRoutingDataSource 类，可以根据用户自定义的规则选择当前的数据源，这样我们可以在执行 SQL 前来决定使用哪个数据源，实现动态路由。因此以下代码实际上都是围绕 AbstractRoutingDataSource 来实现的

# 1.1 DynamicRoutingDataSource.java

public class DynamicRoutingDataSource extends AbstractRoutingDataSource {

	@Override
	protected Object determineCurrentLookupKey() {
		return DynamicDataSourceContextHolder.get();
	}
}

我们通过实现AbstractRoutingDataSource 的模板方法 determineCurrentLookupKey() 方法来决定使用哪个数据源，AbstractRoutingDataSource 的其它具体实现不做赘述

# 1.2 DynamicDataSourceContextHolder.java

public class DynamicDataSourceContextHolder {
	
	private static final ThreadLocal<DataSourceKey> CURRENT_DATASOURCE = new ThreadLocal<>();
	
    /**
     * 清除当前数据源
     */
    public static void clear() {
        CURRENT_DATASOURCE.remove();
    }

    /**
     * 获取当前使用的数据源
     *
     * @return 当前使用数据源的ID
     */
    public static DataSourceKey get() {
        return CURRENT_DATASOURCE.get();
    }

    /**
     * 设置当前使用的数据源
     *
     * @param value 需要设置的数据源ID
     */
    public static void set(DataSourceKey value) {
        CURRENT_DATASOURCE.set(value);
    }
}

一个线程在同一时刻只使用一个数据源，因此使用 ThreadLocal 来保存当前线程使用的数据源。因此 DynamicDataSourceContextHolder 中的方法皆是围绕 ThreadLocal 进行操作。ThreadLocal 中保存的是数据源的 key 值，可以理解为是数据源的名字

# 1.3 DynamicDataSourceConfiguration.java

@EnableConfigurationProperties(MultipleDataSource.class)
@Import({DynamicPlugin.class})
@Configuration
public class DynamicDataSourceConfiguration {
    private static final Logger log = LoggerFactory.getLogger(DynamicDataSourceConfiguration.class);

    @Autowired
    private MultipleDataSource multipleDataSource;


    /**
     * 核心动态数据源
     *
     * @return 数据源实例
     */
    @Bean
    public DataSource dynamicDataSource() {
        // 动态数据源
        DynamicRoutingDataSource dataSource = new DynamicRoutingDataSource();

        // 创建读写数据源 可根据需求来做改变
        DataSource dbWrite = database(DataSourceKey.WRITE);
        DataSource dbRead = database(DataSourceKey.READ);
        // 默认数据源
        dataSource.setDefaultTargetDataSource(dbWrite);

        // 将创建的数据源按照 key value 的形式保存到动态数据源中
        // key 是 DynamicDataSourceContextHolder 中 ThreadLocal 保存的值
        Map<Object, Object> dataSourceMap = new HashMap<>(4);
        dataSourceMap.put(DataSourceKey.READ, dbRead);
        dataSourceMap.put(DataSourceKey.WRITE, dbWrite);
        dataSource.setTargetDataSources(dataSourceMap);
        return dataSource;
    }

    public DataSource database( DataSourceKey key) {
        log.debug("{}库已连接", key);
        return multipleDataSource.getDataSources().get(key.name().toLowerCase());
    }
}

配置类，真正配置数据源的地方

1.4 其它代码

DataSourceKey.java

public enum DataSourceKey {
	READ,
	WRITE
}

数据源对应的 key 值，根据需求自行配置

MultipleDataSource.java

@ConfigurationProperties(prefix = "multiple")
public class MultipleDataSource {

    private Map<String, HikariDataSource> dataSources;

    public Map<String, HikariDataSource> getDataSources() {
        return dataSources;
    }

    public void setDataSources(Map<String, HikariDataSource> dataSources) {
        this.dataSources = dataSources;
    }
}

配置属性的类，在上面的配置类上加上 @EnableConfigurationProperties(MultipleDataSource.class) 注解，在application.yml 中写动态数据源配置的时候 IDEA 就能够给出提示

application.yml 样例：

multiple:
  data-sources:
    write:
      jdbc-url: jdbc:mysql://localhost:3306/write?allowMultiQueries=true&useUnicode=true&characterEncoding=UTF-8&useSSL=false&serverTimezone=GMT%2B8
      driver-class-name: com.mysql.cj.jdbc.Driver
      username: root
      password: 123456
      
    read:
      driver-class-name: com.mysql.cj.jdbc.Driver
      jdbc-url: jdbc:mysql://localhost:3306/read?serverTimezone=Asia/Shanghai&useUnicode=true&characterEncoding=utf8&allowMultiQueries=true&useSSL=false
      username: root
      password: 123456

二、Mybatis插件

动态数据源的核心代码已经实现了，剩下的就应该考虑在何时切换数据源

网上更多的方式是使用 注解+AOP，优点是可以手动指定数据源，缺点就是麻烦，容易忘记写注解

因此这里提供使用 Mybatis 插件机制实现数据源切换

DyamicPlugin.java

@Intercepts({
        @Signature(type = Executor.class, method = "query", args = {
                MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}),
        @Signature(type = Executor.class, method = "query", args = {
                MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "update", args = {
                MappedStatement.class, Object.class})
})
@Component
public class DynamicPlugin implements Interceptor {
    private static final Logger log = LoggerFactory.getLogger(DynamicPlugin.class);


    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        // 是否存在事务
        boolean synchronizationActive = TransactionSynchronizationManager.isSynchronizationActive();
        // 拦截参数
        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];

        // 如果存在事务 则使用写库
        // 如果不存在事务 则采用读写分离策略
        if (!synchronizationActive) {
            if (SqlCommandType.SELECT.equals(ms.getSqlCommandType())) {
                log.debug("【mybatis拦截】read");
                DynamicDataSourceContextHolder.set(DataSourceKey.READ);
            } else {
                log.debug("【mybatis拦截】write");
                DynamicDataSourceContextHolder.set(DataSourceKey.WRITE);
            }
        } else {
            log.debug("【mybatis拦截】write");
            DynamicDataSourceContextHolder.set(DataSourceKey.WRITE);
        }
        Object res = invocation.proceed();
        // 清空路由选择状态，否则可能会出现先读再写，导致出现在读库写的问题
        DynamicDataSourceContextHolder.clear();
        return res;
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof Executor) {
            return Plugin.wrap(target, this);
        }
        return target;
    }

    @Override
    public void setProperties(Properties properties) {
    }

}

这里的逻辑较为简单，可以根据实际需求进行修改。

---

自己测试暂没发现问题，懒得放测试用例