Spring实现多数据源动态切换

背景

随着业务的发展，数据库压力的增大，如何分割数据库的读写压力是我们需要考虑的问题，而能够动态的切换数据源就是我们的首要目标。

---

基础

Spring作为我们项目的应用容器，也对这方面提供了很好的支持，当我们的持久化框架需要数据库连接时，我们需要做到动态的切换数据源，这些Spring的AbstractRoutingDataSource都给我们留了拓展的空间，可以先来看看抽象类AbstractRoutingDataSource在获取数据库连接时做了什么

private Map resolvedDataSources; //从配置文件读取到的DataSources的Map
 
private DataSource resolvedDefaultDataSource; //默认数据源
  
public Connection getConnection() throws SQLException {
   return determineTargetDataSource().getConnection();
}
 
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 abstract Object determineCurrentLookupKey();

可以看到AbstractRoutingDataSource在决定目标数据源的时候，会先调用determineCurrentLookupKey()方法得到一个key，我们通过这个key从配置好的resolvedDataSources(Map结构)拿到这次调用对应的数据源，而determineCurrentLookupKey()开放出来让我们实现

---

实现

前面提到我们可以通过实现AbstractRoutingDataSource的determineCurrentLookupKey()方法来决定这次调用的数据源。首先说一下思路：当我们的一个线程需要针对数据库做一系列操作时，每次都会去调用getConnection()方法获取数据库连接，然后执行完后再释放或归还数据库连接(SqlSessionTemplate就是这么做的)，那么很明显，我们需要能够保证每次调用Dao层方法时都能动态的切换数据源，这就需要Spring的AOP：我们定义一个切面，当我们调用Dao层方法时，执行我们的逻辑来判断这次调用的数据源，AOP切面定义如下：

image.png

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.PACKAGE, ElementType.TYPE, ElementType.METHOD})
public @interface DataSource {
    String value();
}
  
public Object doAround(ProceedingJoinPoint jp) throws Throwable {
 
    Signature signature = jp.getSignature();
 
    String dataSourceKey = getDataSourceKey(signature);
 
    if (StringUtils.hasText(dataSourceKey)) {
        MyDataSource.setDataSourceKey(dataSourceKey);
    }
 
    Object jpVal = jp.proceed();
 
    return jpVal;
}
  
public String getDataSourceKey(Signature signature) {
    if (signature == null) return null;
 
    if (signature instanceof MethodSignature) {
        MethodSignature methodSignature = (MethodSignature) signature;
        Method method = methodSignature.getMethod();
        if (method.isAnnotationPresent(DataSource.class)) {
            return 方法的注解值;
        }
 
        Class declaringClazz = method.getDeclaringClass();
        if (declaringClazz.isAnnotationPresent(DataSource.class)) {
            return 类级别的注解值;
        }
 
        Package pkg = declaringClazz.getPackage();
        return 该包路径的默认数据源;
    }
 
    return null;
}

这里我们就可以得到我们需要的数据源，现在就是如何保存这个值了，因为这个值是我们这个线程才需要使用的，所以综合考虑声明一个ThreadLocal来保存不同线程的不同值，目前已经解决了当前调用方法的数据源和数据源值的保存了，那么回到前面AbstractRoutingDataSource的determineTargetDataSource()方法中，我们就可以重写抽象方法determineCurrentLookupKey，返回我们刚刚保存的数据源的值，代码如下：

public class MyDataSource extends AbstractRoutingDataSource {
 
    private static final ThreadLocal dataSourceKey = new ThreadLocal();
 
    public static void setDataSourceKey(String dataSource) {
        dataSourceKey.set(dataSource);
    }
 
    protected Object determineCurrentLookupKey() {
        String dsName = dataSourceKey.get();
        dataSourceKey.remove(); //这里需要注意的时，每次我们返回当前数据源的值得时候都需要移除ThreadLocal的值，这是为了避免同一线程上一次方法调用对之后调用的影响
        return dsName;
    }
 
}

---

总结

大体的Spring实现多数据源的动态切换思路如上