SpringBoot多數據源

前言

多數據源的核心就是向 IOC 容器注入 AbstractRoutingDataSource 和如何切換數據源。注入的方式可以是注冊 BeanDefinition 或者是構建好的 Bean，切換數據源的方式可以是方法參數或者是注解切換（其他的沒想象出來），具體由需求決定。

我的需求是統計多個庫的數據，將結果寫入另一個數據庫，統計的數據庫數量是不定的，無法通過 @Bean 直接注入，又是統計任務，DAO 層注解切換無法滿足，因此選擇注冊（AbstractRoutingDataSource 的）BeanDefinition 來實現。下面以統計統計中日韓用戶到結果庫為例。

配置文件

master 為結果庫，其他為被統計的數據庫（china、japan 可以用枚舉唯一標識，當然也可以用 String）：

dynamic:
  dataSources:
    master:
      driver-class-name: com.mysql.cj.jdbc.Driver
      url: jdbc:mysql://localhost:3306/result?useUnicode=true&characterEncoding=utf8xxxxxxxx
      username: root
      password: 123456
    china:
      driver-class-name: com.mysql.cj.jdbc.Driver
      url: jdbc:mysql://localhost:3306/china?useUnicode=true&characterEncoding=utf8xxxxxxxx
      username: root
      password: 123456
    japan:
      driver-class-name: com.mysql.cj.jdbc.Driver
      url: jdbc:mysql://localhost:3306/japan?useUnicode=true&characterEncoding=utf8xxxxxxxx
      username: root
      password: 123456
    korea:
      driver-class-name: com.mysql.cj.jdbc.Driver
      url: jdbc:mysql://localhost:3306/korea?useUnicode=true&characterEncoding=utf8xxxxxxxx
      username: root
      password: 123456

對應的配置類：

package com.statistics.dynamicds.core.config;

import com.statistics.dynamicds.core.Country;
import lombok.Data;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Configuration;

import java.util.Map;

import static com.statistics.dynamicds.core.config.MultiDataSourceProperties.PREFIX;

@Data
@Configuration
@ConfigurationProperties(prefix = PREFIX)
public class MultiDataSourceProperties {
  public static final String PREFIX = "dynamic";
  private Map<Country, DataSourceProperties> dataSources;

  @Data
  public static class DataSourceProperties {
    private String driverClassName;
    private String url;
    private String username;
    private String password;
  }
}

package com.statistics.dynamicds.core;

public enum Country {
  MASTER("master", 0),

  CHINA("china", 86),
  JAPAN("japan", 81),
  KOREA("korea", 82),
  // 其他國家省略

  private final String name;
  private final int id;

  Country(String name, int id) {
    this.name = name;
    this.id = id;
  }

  public int getId() {
    return id;
  }

  public String getName() {
    return name;
  }
}

依賴

ORM 用的 JPA，SpringBoot 版本為 2.3.7.RELEASE，通過 Lombok 簡化 GetSet。

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-jpa</artifactId>
</dependency>

<dependency>
     <groupId>org.projectlombok</groupId>
     <artifactId>lombok</artifactId>
     <version>1.18.22</version>
     <scope>provided</scope>
</dependency>

構建 AbstractRoutingDataSource

Spring 的動態數據源需要注入 AbstractRoutingDataSource，因為配置文件中被統計數據源不是固定的，所以不能通過 @Bean 注解注入，需要手動構建。

要在啟動類加上 @Import(MultiDataSourceImportBeanDefinitionRegistrar.class)。

要在啟動類加上 @Import(MultiDataSourceImportBeanDefinitionRegistrar.class)。

要在啟動類加上 @Import(MultiDataSourceImportBeanDefinitionRegistrar.class)，重要的事情寫三行。

package com.statistics.dynamicds.autoconfig;

import com.statistics.dynamicds.core.DynamicDataSourceRouter;
import com.statistics.dynamicds.core.Country;
import com.statistics.dynamicds.core.config.MultiDataSourceProperties;
import com.zaxxer.hikari.HikariDataSource;
import org.springframework.beans.factory.support.AbstractBeanDefinition;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.boot.context.properties.bind.Binder;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.env.Environment;
import org.springframework.core.type.AnnotationMetadata;

import javax.annotation.Nonnull;
import java.util.Map;
import java.util.stream.Collectors;

import static com.statistics.dynamicds.core.config.MultiDataSourceProperties.PREFIX;

public class MultiDataSourceImportBeanDefinitionRegistrar implements ImportBeanDefinitionRegistrar, EnvironmentAware {
  public static final String DATASOURCE_BEANNAME = "dynamicDataSourceRouter";
  private Environment environment;

  @Override
  public void registerBeanDefinitions(@Nonnull AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry) {
    MultiDataSourceProperties multiDataSourceProperties = Binder.get(environment)
            .bind(PREFIX, MultiDataSourceProperties.class)
            .orElseThrow(() -> new RuntimeException("no found dynamicds config"));
    final HikariDataSource[] defaultTargetDataSource = {null};
    Map<Country, HikariDataSource> targetDataSources = multiDataSourceProperties.getDataSources().entrySet().stream()
            .collect(Collectors.toMap(
                    Map.Entry::getKey,
                    entry -> {
              MultiDataSourceProperties.DataSourceProperties dataSourceProperties = entry.getValue();
              HikariDataSource dataSource = DataSourceBuilder.create()
                      .type(HikariDataSource.class)
                      .driverClassName(dataSourceProperties.getDriverClassName())
                      .url(dataSourceProperties.getUrl())
                      .username(dataSourceProperties.getUsername())
                      .password(dataSourceProperties.getPassword())
                      .build();
              dataSource.setPoolName("HikariPool-" + entry.getKey());
              if (Country.MASTER == entry.getKey()) {
                defaultTargetDataSource[0] = dataSource;
              }
              return dataSource;
            }));
    targetDataSources.remove(Country.MASTER);
    AbstractBeanDefinition beanDefinition = BeanDefinitionBuilder.genericBeanDefinition(DynamicDataSourceRouter.class)
            .addConstructorArgValue(defaultTargetDataSource[0])
            .addConstructorArgValue(targetDataSources)
            .getBeanDefinition();
    registry.registerBeanDefinition(DATASOURCE_BEANNAME, beanDefinition);
  }

  @Override
  public void setEnvironment(@Nonnull Environment environment) {
    this.environment = environment;
  }
}

上面代碼中 MultiDataSourceProperties 不是由 @Resource 或者 @Autowired 獲取的是因為 ImportBeanDefinitionRegistrar 執行的很早，此時 @ConfigurationProperties 的配置參數類還沒有注入，因此要手動獲取（加 @ConfigurationProperties 注解是為了使 IOC 容器中其他 Bean 能獲取配置的 Country，以此來切換數據源）。

下面是 AbstractRoutingDataSource 的實現類 DynamicDataSourceRouter：

package com.statistics.dynamicds.core;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.Map;

public class DynamicDataSourceRouter extends AbstractRoutingDataSource {
  public DynamicDataSourceRouter(Object defaultTargetDataSource, Map<Object, Object> targetDataSources) {
    this.setDefaultTargetDataSource(defaultTargetDataSource);
    this.setTargetDataSources(targetDataSources);
  }

  @Override
  protected Object determineCurrentLookupKey() {
    return DataSourceContextHolder.getLookupKey();
  }
}

數據源切換

數據源的切換由 DataSourceContextHolder 控制：

package com.statistics.dynamicds.core;

public class DataSourceContextHolder {
  private static final ThreadLocal<Country> HOLDER = ThreadLocal.withInitial(() -> Country.MASTER);

  public static void setLookupKey(Country lookUpKey) {
    HOLDER.set(lookUpKey);
  }

  public static Country getLookupKey() {
    return HOLDER.get();
  }

  public static void clear() {
    HOLDER.remove();
  }
}

事務

當沒有事務時可以隨意切換數據源。如果有事務就比較麻煩了，必須在執行 @Transactional 注解修飾的方法前（事務開始前）完成數據源的切換，也就是必須通過 DataSourceContextHolder.setLookupKey() 手動切換數據源，且被 @Transactional 修飾的方法內不能有不同數據源的寫（但可以用 @Transactional(propagation = Propagation.NOT_SUPPORTED) 修飾 DAO 方法完成其他數據源的讀）。

目錄

.
└─com
    └─statistics
        │  StatisticsApplication.java
        │
        ├─dao
        │      UserDao.java
        │
        ├─dynamicds
        │  ├─autoconfig
        │  │      MultiDataSourceImportBeanDefinitionRegistrar.java
        │  │
        │  └─core
        │      │  DataSourceContextHolder.java
        │      │  DynamicDataSourceRouter.java
        │      │  Country.java
        │      │
        │      └─config
        │              MultiDataSourceProperties.java

總結

以上就完成了多數據源配置，使用時只需要調用 DataSourceContextHolder.setLookupKey() 就可以切換了。

如果無法用枚舉標識數據源也可以換成 String，關於這個數據源的其他信息在內部類 DataSourceProperties 加一個 map 即可，總之就是按照自己的需求擴展。