Spring中实现动态数据源切换，基于AbstractRoutingDataSource

背景

在项目开发过程中，我们可能会遇到一个场景：某个类型数据源有多个数据源实例，需要我们按照不同的请求切换到不同数据源去。
而目前绝大多数java应用都是基于Spring框架来开发，我们很多时候相关的数据源连接都是交给了Spring框架去管理，这就需要Spring能够支持动态数据源切换。

方案

Spring中预留了这个接口，通过AbstractRoutingDataSource能够动态切换数据源。

public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {

这是一个抽象类，预留了一个抽象方法：

protected abstract Object determineCurrentLookupKey();

我们知道，数据源一般都会提供一个getConnection方法来获取一个连接，在AbstractRoutingDataSource 实现如下：

	@Override
	public Connection getConnection() throws SQLException {
		return determineTargetDataSource().getConnection();
	}
	protected DataSource determineTargetDataSource() {
		Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
		Object lookupKey = determineCurrentLookupKey();
		DataSource dataSource = this.resolvedDataSources.get(lookupKey);
		if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
			dataSource = this.resolvedDefaultDataSource;
		}
		if (dataSource == null) {
			throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
		}
		return dataSource;
	}

可以看到，AbstractRoutingDataSource 获取连接的主要逻辑就是通过determineCurrentLookupKey获取到一个数据源的关联key，然后从resolvedDataSources中去获取。
而resolvedDataSources的初始化，则放在afterPropertiesSet中：

	@Override
	public void afterPropertiesSet() {
		if (this.targetDataSources == null) {
			throw new IllegalArgumentException("Property 'targetDataSources' is required");
		}
		this.resolvedDataSources = CollectionUtils.newHashMap(this.targetDataSources.size());
		this.targetDataSources.forEach((key, value) -> {
			Object lookupKey = resolveSpecifiedLookupKey(key);
			DataSource dataSource = resolveSpecifiedDataSource(value);
			this.resolvedDataSources.put(lookupKey, dataSource);
		});
		if (this.defaultTargetDataSource != null) {
			this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
		}
	}

这里起始就是通过targetDataSources中指定的数据源复制到resolvedDataSources 中去。因此如果多数源是固定的，那么只需要实现determineCurrentLookupKey方法即可。但是如果多数据源不固定，比如可能会有数据源的变更，那么这种实现是不能够支持，因为这种实现从服务启动的视乎，后续数据源就不能发生变更，这需要我们自己实现determineTargetDataSource.
下面是一个参考实现：

public class DataSourceContextHolder {

    private static final ThreadLocal<String> DATASOURCE_CONTEXT_KEY_HOLDER = new ThreadLocal<>();


    public static void switchDataSource(String key){
        log.info("Switch to data source:" + key);
        DATASOURCE_CONTEXT_KEY_HOLDER.set(key);
    }
    public static String getDataSourceKey(){
        return DATASOURCE_CONTEXT_KEY_HOLDER.get() ;
    }

}

public class DynamicDataSource extends AbstractRoutingDataSource {

    private Map<Object, Object> targetDataSources = new HashMap<>();
    private Map<Object, DataSource> dataSources = new HashMap<>();
    public DynamicDataSource (){
        super.setDefaultTargetDataSource(null);
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }
 @Override
    protected DataSource determineTargetDataSource() {
        Object dataSourceKey = determineCurrentLookupKey();
        return dataSources.get(dataSourceKey);
    }
    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceContextHolder .getDataSourceKey();
    }
    public synchronized void addDataSource(String key, DataSource dataSource){
        targetDataSources.put(key,dataSource);
        dataSources.put(key,dataSource);
        log.info("add dynamic dataSource for contextKey= {} ",key);
    }
}

这样我们通过DataSourceContextHolder 来调整当前线程关联的数据源。

上面这种处理在单个请求上下文中你需要对动态数据源调整场景下无法实现，因为Srping基于事务处理下，在一次请求上线文中（也就是同一个线程中），对于同一个数据源，一旦获取到实际的连接（Connection）之后，会保存到ThreadLocal中，之后不会再去获取连接，而是从ThreadLocal中获取，这也是Spring事务实现的一个基础原理，上述实现只能在请求开始时数据源没确定的时候生效，一旦连接获取到之后，之后实际是不会再进行切换的。

遇到这种情况时，我们对AbstractRoutingDataSource 获取到的Connection进行包装，使用装饰器模式进行额外的处理，实现一个DynamicConnection。

@Slf4j
public class DynamicConnection implements Connection {
    private Connection directConnection;
    private String contextKey;
	public void closeDirectConnection(){
        if (this.directConnection != null){
            log.info("close direct connection for contextKey={}",contextKey);
            try {
                this.directConnection.close();
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }
        @Override
    public void close() throws SQLException {
        log.info("close dynamic connection and remove from current context");
        directConnection.close();
    }
    // 其他方法直接使用directConnection的对应方法实现即可
}

    @Override
    public Connection getConnection() throws SQLException {
        Connection directConnection =  getDirectConnection();
        String contextKey= DynamicDataSourceContext.getDataSourceKey();
        DynamicConnection dynamicConnection= new DynamicConnection ();
        dynamicConnection.setDirectConnection(directConnection);
        dynamicConnection.setContextKey(contextKey);
        log.info("getConnection from wrapper connection with contextKey= {}",contextKey;
        // 将连接放入到本次请求的上下文中去
      DynamicDataSourceContext.setCurrentConnection(dynamicConnection);
        return dynamicConnection;
    }

当发生数据源切换时，通过AOP触发，或者直接调用该方法，切换DynamicConnection 中的directConnection属性

private void switchConnection(String contextKey) {
        DynamicDataSourceContext.switchDataSource(contextKey);
        if (DynamicDataSourceContext.getCurrentConnection() != null) {
            try {
                log.info("switchDirectConnection to {}",contextKey);
                DynamicDataSourceContext.getCurrentConnection().closeDirectConnection();
                 DynamicDataSource dynamicDataSource = SpringContextUtil.getBean(DynamicDataSource .class);
                DynamicDataSourceContext.getCurrentConnection().setDirectConnection(dynamicDataSource .getDirectConnection());
                DynamicDataSourceContext.getCurrentConnection().setContextKey(contextKey);
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }

这样我们就在一次请求中能够动态切换连接。但是这种场景下，需要确保事务不会跨多个数据源。