Spring数据源动态切换

Spring数据源动态切换

利用spring-jdbc包中的AbstractRoutingDataSource类，从名字就可以看出来这是个抽象类。

先来看看这个类的相关类图

spring动态数据源类图.png

可以看到这个类实现了DataSource接口，这意味着其继承类可以在任何需要DataSource的地方使用。

DataSource接口中只定义了两个方法

Connection getConnection() throws SQLException;

Connection getConnection(String username, String password)
    throws SQLException;

所以我们重点关注AbstractRoutingDataSource中如何实现这两个方法。

    @Override
    public Connection getConnection() throws SQLException {
        return determineTargetDataSource().getConnection();
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return determineTargetDataSource().getConnection(username, password);
    }

继续追踪

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;
    }

这个方法被设计成protected，说明子类可以重写。该方法中调用了AbstractRoutingDataSource中唯一一个抽象方法

/**
     * Determine the current lookup key. This will typically be
     * implemented to check a thread-bound transaction context.
     * 
Allows for arbitrary keys. The returned key needs
     * to match the stored lookup key type, as resolved by the
     * {@link #resolveSpecifiedLookupKey} method.
     */
    protected abstract Object determineCurrentLookupKey();

这个方法便是我们实现动态数据源切换的关键了。

determineTargetDataSource中可以看到最终使用的DataSource对象是从resolvedDataSources中获取的。其定义如下：

private Map resolvedDataSources;

那么这个map是什么时候填充的呢？

   @Override
    public void afterPropertiesSet() {
        if (this.targetDataSources == null) {
            throw new IllegalArgumentException("Property 'targetDataSources' is required");
        }
        this.resolvedDataSources = new HashMap(this.targetDataSources.size());
        for (Map.Entry entry : this.targetDataSources.entrySet()) {
            Object lookupKey = resolveSpecifiedLookupKey(entry.getKey());
            DataSource dataSource = resolveSpecifiedDataSource(entry.getValue());
            this.resolvedDataSources.put(lookupKey, dataSource);
        }
        if (this.defaultTargetDataSource != null) {
            this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
        }
    }

可以清晰的看到是用targetDataSources里面的数据填充的。

接下来就是最终的解决方案了。

1. 定义数据源

public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
    
        String datasource = ContextHolder.getCustomerType();
        return datasource;
    }
}

其中test1,test2以及default都是我们定义的DataSource bean。

2. 动态数据源操作

public class ContextHolder {


  public static final String DATA_SOURCE_TEST1 = "test1";

  /**
   * 猜你喜欢数据库
   */
  public static final String DATA_SOURCE_TEST2 = "test2";

  /**
   * 当前系统的数据库
   */
  public static final String DATA_SOURCE_DEFAULT = "default";

  // 利用ThreadLocal解决线程安全问题
  private static final ThreadLocal contextHolder = new ThreadLocal();

  public static void setCustomerType(String customerType) {

    contextHolder.set(customerType);
  }

  public static String getCustomerType() {

    return contextHolder.get();
  }

  public static void clearCustomerType() {

    contextHolder.remove();
  }
}

3. 定义切面

public class DataSourceInterceptor {

    public void setDefaultDataSource() {
        ContextHolder.setCustomerType(ContextHolder.DATA_SOURCE_DEFAULT);
    }

    public void setTest1DataSource() {
        ContextHolder.setCustomerType(ContextHolder.DATA_SOURCE_TEST1);
    }

    public void setTest2DataSource() {
        ContextHolder.setCustomerType(ContextHolder.DATA_SOURCE_TEST2);
    }
}

这样在调用相应的方法时就能自动切换数据源了，但是设置完毕数据源后并没有清除，这意味着需要手动清除。可以加上aop after在每个方法调用后清除掉当前的数据源。