AbstractRoutingDataSource实现动态数据源切换

young 680 2021-11-24

AbstractRoutingDataSource

spring-jdbc的包中，提供了AbstractRoutingDataSource用于数据源路由操作

public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {

该类实现了InitializingBean接口，说明初始化时会调用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属性不能为空，需要给这个属性赋值

查看resolveSpecifiedDataSource方法

protected DataSource resolveSpecifiedDataSource(Object dataSource) throws IllegalArgumentException {
   if (dataSource instanceof DataSource) {
      return (DataSource) dataSource;
   }
   else if (dataSource instanceof String) {
      return this.dataSourceLookup.getDataSource((String) dataSource);
   }
   else {
      throw new IllegalArgumentException(
            "Illegal data source value - only [javax.sql.DataSource] and String supported: " + dataSource);
   }
}

如果dataSource是字符串，会生成一个DataSource，由此可见，如果数据源连接使用的是jndi，可以直接将jndi的名称传入

所以这个初始化的操作实际就是将targetDataSources中value，处理为DataSource，并按我们存入的映射关系进行映射，默认数据源不为空时，将其处理成一个DataSource

该抽象类拥有一个唯一的抽象方法

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

在determineTargetDataSource方法中调用了该抽象方法

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

由此可见该方法是获取一个key，通过key来获取需要的数据源，如果找不到对应的数据源，则返回设置的默认数据源

添加数据源配置

我们先创建一个配置类继承AbstractRoutingDataSource

@Configuration
public class DynamicDataSource extends AbstractRoutingDataSource {
    

    @Override
    protected Object determineCurrentLookupKey() {
        return null;
    }
}

获取DataSource的操作应该是线程隔离的，所以使用ThreadLocal进行key的获取

AbstractRoutingDataSource中已经存在了初始化方法，所以我们重写afterPropertiesSet方法，进行targetDataSources和defaultTargetDataSource的赋值操作

创建DBContextManage来操作TheadLocal

public class DBContextManage {
    private DBContextManage(){}

    private static final ThreadLocal<String> DB_CONTEXT = new ThreadLocal<>();
	// 设置key
    public static void set(String dbName){
        DB_CONTEXT.set(dbName);
    }
	// 获取key
    public static String get(){
       return DB_CONTEXT.get();
    }
    // 清除ey
    public static void clear(){
        DB_CONTEXT.remove();
    }
}

可以基于数据库记录或配置文件来获取DataSource的配置

基于数据库

基于数据库时，需用@Primary指定默认DataSource，然后通过数据库查询获得其他数据源的链接信息

@Configuration
public class DynamicDataSource extends AbstractRoutingDataSource {
    @Resource
    private DbDao dbDao;

    @Override
    public void afterPropertiesSet() {
        List<Db> list = dbDao.selectAll();
        Map<Object, Object> targetDataSources = list.stream().collect(Collectors.toMap(e -> e.getDbId(), e -> e.getJndi()));
        setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }

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

基于配置文件

datasource:
  configs:
    - jdbcUrl: jdbc:oracle:thin:@10.10.200.42:1521:orcl
      driverClassName: oracle.jdbc.OracleDriver
      userName: aaaaaaa
      password: aaaaaaa
      key: db1
    - jdbcUrl: jdbc:oracle:thin:@10.10.200.42:1521:orcl
      driverClassName: oracle.jdbc.OracleDriver
      userName: bbbbbbbb
      password: bbbbbbbb
      key: db2

配置对应类

@Data
@Configuration
@ConfigurationProperties(prefix = "datasource")
public class DynamicDataSourceConfig {
    private List<Config> configs = Collections.emptyList();

    @Data
    public static class Config{
        private String jdbcUrl;
        private String driverClassName;
        private String userName;
        private String password;
        private String key;
    }
}

数据源配置

@Configuration
public class DynamicDataSource extends AbstractRoutingDataSource {
   @Autowired
   private DynamicDataSourceConfig dynamicDataSourceConfig;

   @Override
   public void afterPropertiesSet() {
      Map<Object, Object> targetDataSourcesMap = dynamicDataSourceConfig.getConfigs().stream()
            .collect(Collectors.toMap(DynamicDataSourceConfig.Config::getKey,
                  e -> DataSourceBuilder.create().driverClassName(e.getDriverClassName()).url(e.getJdbcUrl())
                        .username(e.getUserName()).password(e.getPassword()).build()));
      super.setTargetDataSources(targetDataSourcesMap);
      super.afterPropertiesSet();
   }

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

测试

创建一个测试Dao

public interface TestDao {    List<String> test();}<select id="test" resultType="string">    select distinct applied_by from changelog</select>

创建一个测试Controller

@RestController
public class TestController {

    @Autowired
    private TestDao testDao;

    @RequestMapping("/test")
    public void test(){
        System.out.println("====>>dynamic datasource test");
        DBContextManage.set("db1");
        System.out.println(testDao.test());
        DBContextManage.set("db2");
        System.out.println(testDao.test());
    }
}

ControllerTestLog.png

由此可见数据源的切换已经成功了

但是我们开发的过程中不会这么操作，一般采用自定义注解加AOP的方式来实现

基于注解加AOP进行数据源切换

创建自定义注解

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface DBRoute {
    String value();
}

创建AOP

@EnableAspectJAutoProxy@Configuration@Aspectpublic class DBRouteAspectJ {    @Pointcut("@annotation(com.example.dynamicdatasource.DBRoute)")    public void pointCut() {}    @Around(value = "pointCut()&&@annotation(annotation)")    public Object dbRoute(ProceedingJoinPoint joinPoint,DBRoute annotation) throws Throwable{        // 设置key        DBContextManage.set(annotation.value());        // 执行方法        Object proceed = joinPoint.proceed();        // 清理ThreadLocal，避免内存泄露        DBContextManage.clear();        return proceed;    }}

测试

创建测试Service

@Service
public class TestService {
    @Autowired
    private TestDao testDao;

    @DBRoute("db1")
    public void test1(){
        System.out.println("======>test1");
        System.out.println(testDao.test());
    }

    @DBRoute("db2")
    public void test2(){
        System.out.println("======>test2");
        System.out.println(testDao.test());
    }
}

修改测试Controller

@RestController
public class TestController {

    @Autowired
    private TestService testService;

    @RequestMapping("/test")
    public void test(){
        System.out.println("====>>dynamic datasource test");
        testService.test1();
        testService.test2();
    }
}

AopTestLog.png

由日志可见，基于AOP进行数据源切换成功了

目前所有的操作都是读操作，不考虑事务的

springboot mybatis