Spring Boot JPA 事务中动态切换数据源

网上通过 RoutingDataSource + ThreadLocal + AOP 实现动态切换数据源的文章很多，但是一旦加上@Transactional就无法切换了。原因是事务提交时才会调用AbstractRoutingDataSource的determineCurrentLookupKey方法， 获取当前数据源。而在事务中就算切换多次数据源，只会使用事务提交时的当前数据源。因此，要在事务中切换数据源，必须使用@Transactional(propagation = Propagation.REQUIRES_NEW)，开启新的事务。关键代码如下：

• RoutingDataSourceContext
  通过ThreadLocal上下文管理当前数据源，为了防止内存泄漏，实现AutoCloseable接口，使用try-with-resources，使用完毕后自动上下文内容。

public class RoutingDataSourceContext implements AutoCloseable {

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

    public static String getDataSourceRoutingKey() {
        return threadLocalDataSourceKey.get();
    }

    public RoutingDataSourceContext(String key) {
        threadLocalDataSourceKey.set(key);
    }

    @Override
    public void close() throws Exception {
        threadLocalDataSourceKey.remove();
    }

}

• RoutingDataSource
  实现determineCurrentLookupKey方法，从上下文获取动态切换的数据源名称。

@Slf4j
public class RoutingDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        String key = RoutingDataSourceContext.getDataSourceRoutingKey();
        log.debug("查询当前数据源名称为{}", key);
        return key;
    }

}

• DataSourceConfig
  从配置文件中读取多个数据源，当然也可以在代码中自动注入RoutingDataSource，添加或删除数据源。

@Configuration
@ConfigurationProperties("ba.datasource")
public class DataSourceConfig {

    @Getter
    @Setter
    private DataSourceProperties master;

    @Getter
    @Setter
    private Map<String, DataSourceProperties> slaves = new HashMap<>();

    @Bean
    DataSource dataSource() {
        RoutingDataSource dataSource = new RoutingDataSource();

        dataSource.setDefaultTargetDataSource(master.initializeDataSourceBuilder().build());

        Map<Object, Object> targetDataSources = new HashMap<>();
        slaves.forEach((key, slave) -> {
            targetDataSources.put(key, slave.initializeDataSourceBuilder().build());
        });
        dataSource.setTargetDataSources(targetDataSources);

        return dataSource;
    }

}

• application.yml
  配置数据源属性。

ba:
  datasource:
    master:
      driver-class-name: org.postgresql.Driver
      url: jdbc:postgresql://localhost:5432/master
      username: postgres
      password: 123456
    slaves:
      slave1:
        driver-class-name: org.postgresql.Driver
        url: jdbc:postgresql://localhost:5432/slave1
        username: postgres
        password: 123456

• RoutingWith
  定义注解，通过在方法上添加注解，声明将要切换的数据源名称。

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RoutingWith {

    String value();

}

• RoutingAspect
  通过AOP，拦截使用了@RoutingWith的方法，动态切换上下文中的数据源名称。注意添加@Order(-1)，使得在事务的AOP外层拦截。

@Slf4j
@Aspect
@Component
@Order(-1)
public class RoutingAspect {

    @Around("@annotation(routingWith)")
    public Object routingWithDataSource(ProceedingJoinPoint joinPoint, RoutingWith routingWith) throws Throwable {
        String key = routingWith.value();
        log.debug("切换数据源为{}", key);
        try (RoutingDataSourceContext ctx = new RoutingDataSourceContext(key)) {
            return joinPoint.proceed();
        }
    }

}

• UserManager
  使用实例。addUser方法添加了@Transactional，如果test2方法不添加@Transactional(propagation = Propagation.REQUIRES_NEW)，将会保存两条记录到默认的master数据源。添加注解后，test1的数据还是保存在默认的master数据源，而test2的数据将会保存到声明的slave1数据源中。

@Component
public class UserManager {

    @Autowired
    private UserDAO userDAO;

    @Autowired(required = false)
    private UserManager userManager;

    @Transactional
    public void addUser() {
        userManager.test1();

        userManager.test2();
    }

    public void test1() {
        UserDVO user = new UserDVO();
        user.setUsername("Jack Ma");
        user.setPassword("123456");
        userDAO.save(user);
    }

    @Transactional(propagation = Propagation.REQUIRES_NEW)
    @RoutingWith("slave1")
    public void test2() {
        UserDVO user = new UserDVO();
        user.setUsername("Zhouyi Ma");
        user.setPassword("654321");
        userDAO.save(user);
    }

}