SpringBoot多數據源自動切換

現在的企業服務逐漸地呈現出數據的指數級增長趨勢，無論從數據庫的選型還是搭建，大多數的團隊都開始考慮多樣化的數據庫來支撐存儲服務。例如分布式數據庫、Nosql數據庫、內存數據庫、關系型數據庫等等。再到后端開發來說，服務的增多，必定需要考慮到多數據源的切換使用來兼容服務之間的調用。

一、引入依賴

<!-- 核心啟動器 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter</artifactId>
</dependency>
<!-- jdbc 操作數據庫API -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-jdbc</artifactId>
</dependency>
<!-- 數據庫驅動 -->
<dependency>
    <groupId>mysql</groupId>
    <artifactId>mysql-connector-java</artifactId>
    <scope>runtime</scope>
</dependency>
<!-- aop 切面 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-aop</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-test</artifactId>
    <scope>test</scope>
</dependency>

二、application.properties

spring:
  datasource:
    # default數據源, 這里有點郁悶，將如果默認數據源是這樣的形式spring.datasource.default.url,會導致錯誤
    # 根據網上的一些說法啟動時不會開啟自動配置數據庫:@SpringBootApplication(exclude = {DataSourceAutoConfiguration.class})
    # 但是還是不行，最后只能這樣配置了
    url: jdbc:mysql://192.168.178.5:12345/mydb?characterEncoding=UTF-8&useUnicode=true&useSSL=false
    username: root
    password: 123456
    driver: com.mysql.jdbc.Driver
    type: com.zaxxer.hikari.HikariDataSource

    #其余數據源
    names: tb1,tb2
    tb1:
      url: jdbc:mysql://192.168.178.5:12345/mydb2?characterEncoding=UTF-8&useUnicode=true&useSSL=false
      username: root
      password: 123456
      driver:  com.mysql.jdbc.Driver

    tb2:
      url: jdbc:mysql://192.168.178.5:12345/mydb3?characterEncoding=UTF-8&useUnicode=true&useSSL=false
      username: root
      password: 123456
      driver:  com.mysql.jdbc.Driver

三、具體代碼實施

1. 使用ThreadLocal創建一個線程安全的類，存放當前線程的數據源類型

public class DynamicDataSourceContextHolder {

    //存放當前線程使用的數據源類型信息
    private static final ThreadLocal<String> contextHolder = new ThreadLocal<String>();

    //存放數據源id
    public static List<String> dataSourceIds = new ArrayList<String>();

    //設置數據源
    public static void setDataSourceType(String dataSourceType) {
        contextHolder.set(dataSourceType);
    }

    //獲取數據源
    public static String getDataSourceType() {
        return contextHolder.get();
    }

    //清除數據源
    public static void clearDataSourceType() {
        contextHolder.remove();
        System.out.println("清除數據源:" + contextHolder.get());
    }

    //判斷當前數據源是否存在
    public static boolean isContainsDataSource(String dataSourceId) {
        return dataSourceIds.contains(dataSourceId);
    }
}

2. 創建一個DynamicDataSource重寫AbstractRoutingDataSource的determineCurrentLookupKey()方法

/**
 * AbstractRoutingDataSource的內部維護了一個名為targetDataSources的Map，
 * 並提供的setter方法用於設置數據源關鍵字與數據源的關系，實現類被要求實現其determineCurrentLookupKey()方法，
 * 由此方法的返回值決定具體從哪個數據源中獲取連接。
 */
public class DynamicDataSource extends AbstractRoutingDataSource {
    @Override
    protected Object determineCurrentLookupKey() {
        String dataSource = DynamicDataSourceContextHolder.getDataSourceType();
        System.out.println("當前數據源：" + dataSource);
        return dataSource;
    }

}

3. 創建一個Register實現ImportBeanDefinitionRegistrar, EnvironmentAware

/**
 * ImportBeanDefinitionRegistrar介紹
 *  1.ImportBeanDefinitionRegistrar接口不是直接注冊Bean到IOC容器，它的執行時機比較早，
 *      准確的說更像是注冊Bean的定義信息以便后面的Bean的創建。
 *  2. ImportBeanDefinitionRegistrar接口提供了registerBeanDefinitions方便讓子類進
 *      行重寫。該方法提供BeanDefinitionRegistry類型的參數，讓開發者調用BeanDefinitionRegistry的registerBeanDefinition方法傳入BeanDefinitionName和對應的BeanDefinition對象，直接往容器中注冊。
 *  3. ImportBeanDefinitionRegistrar只能通過由其它類import的方式來加載，通常是主啟動類類或者注解。
 *  4. 這里要特別注意：ImportBeanDefinitionRegistrar有兩個重載的registerBeanDefinitions方法，我們只需要重寫其中一個即可否則容易出錯
 *      (1) 如果重寫了兩個方法容易出現這個問題，三個參數的registerBeanDefinitions方法為空邏輯，兩個參數的registerBeanDefinitions方法有實際的
 *          代碼邏輯，這樣會導致代碼邏輯不能實際被執行，需要在三個參數的那個方法再調一個兩個參數的方法
 *      (2) 如果重寫了兩個方法，就只能必須將實際的代碼邏輯寫在有三個參數的registerBeanDefinitions方法中，然后兩個參數的方法不寫任何邏輯即可
 *      通用查看ImportBeanDefinitionRegistrar的實際源碼即可知道問題所在。
 */

public class DynamicDataSourceRegister implements ImportBeanDefinitionRegistrar, EnvironmentAware {

    //指定默認數據源(springboot2.0默認數據源是hikari，大家也可以使用DruidDataSource)
    private static final String DATASOURCE_TYPE_DEFAULT = "com.zaxxer.hikari.HikariDataSource";
    //默認數據源 javax.sql.DataSource
    private DataSource defaultDataSource;
    //用戶自定義數據源
    private Map<String, DataSource> slaveDataSources = new HashMap<>();

    @Override
    public void setEnvironment(Environment environment) {
        this.initDefaultDataSource(environment);
        this.initMutilDataSources(environment);
    }

    private void initDefaultDataSource(Environment env) {
        // 讀取主數據源,解析yml文件
        Map<String, Object> dsMap = new HashMap<>();
        dsMap.put("driver", env.getProperty("spring.datasource.driver"));
        dsMap.put("url", env.getProperty("spring.datasource.url"));
        dsMap.put("username", env.getProperty("spring.datasource.username"));
        dsMap.put("password", env.getProperty("spring.datasource.password"));
        dsMap.put("type", env.getProperty("spring.datasource.type"));
        defaultDataSource = buildDataSource(dsMap);
    }

    private void initMutilDataSources(Environment env) {
        // 讀取配置文件獲取更多數據源
        String dsPrefixs = env.getProperty("spring.datasource.names");
        for (String dsPrefix : dsPrefixs.split(",")) {
            // 多個數據源
            Map<String, Object> dsMap = new HashMap<>();
            dsMap.put("driver", env.getProperty("spring.datasource." + dsPrefix.trim() + ".driver"));
            dsMap.put("url", env.getProperty("spring.datasource." + dsPrefix.trim() + ".url"));
            dsMap.put("username", env.getProperty("spring.datasource." + dsPrefix.trim() + ".username"));
            dsMap.put("password", env.getProperty("spring.datasource." + dsPrefix.trim() + ".password"));
            DataSource ds = buildDataSource(dsMap);
            slaveDataSources.put(dsPrefix, ds);
        }
    }

    private DataSource buildDataSource(Map<String, Object> dataSourceMap) {
        try {
            Object type = dataSourceMap.get("type");
            if (type == null) {
                type = DATASOURCE_TYPE_DEFAULT;// 默認DataSource
            }

            Class<? extends DataSource> dataSourceType;
            dataSourceType = (Class<? extends DataSource>) Class.forName((String) type);
            String driverClassName = dataSourceMap.get("driver").toString();
            String url = dataSourceMap.get("url").toString();
            String username = dataSourceMap.get("username").toString();
            String password = dataSourceMap.get("password").toString();
            // 自定義DataSource配置
            DataSourceBuilder factory = DataSourceBuilder.create().driverClassName(driverClassName).url(url)
                    .username(username).password(password).type(dataSourceType);
            return factory.build();
        } catch (ClassNotFoundException e) {
            System.out.println("找不到指定的類");
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 注入DynamicDataSource的 bean 定義,
     */
    @Override
    public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry) {
        Map<Object, Object> targetDataSources = new HashMap<Object, Object>();
        //添加默認數據源
        targetDataSources.put("default", this.defaultDataSource);
        DynamicDataSourceContextHolder.dataSourceIds.add("default");
        //添加其他數據源
        targetDataSources.putAll(slaveDataSources);
        DynamicDataSourceContextHolder.dataSourceIds.addAll(slaveDataSources.keySet());
        //創建DynamicDataSource
        GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
        beanDefinition.setBeanClass(DynamicDataSource.class);
        beanDefinition.setSynthetic(true);
        MutablePropertyValues mpv = beanDefinition.getPropertyValues();
        //defaultTargetDataSource 和 targetDataSources屬性是 AbstractRoutingDataSource的兩個屬性Map
        mpv.addPropertyValue("defaultTargetDataSource", defaultDataSource);
        mpv.addPropertyValue("targetDataSources", targetDataSources);
        //注冊 - BeanDefinitionRegistry
        registry.registerBeanDefinition("dataSource", beanDefinition);
        System.out.println("Dynamic DataSource Registry");
    }
}

4. 創建注解類@DbName

/**
 * @DbName注解用於類、方法上
 */
@Target({ElementType.TYPE, ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface DbName {
    String value();
}

5. 創建aop切面類ChooseDbAspect

@Component
@Order(-1)//這里一定要保證在@Transactional之前執行
@Aspect
public class ChooseDbAspect {
    /**
     * 切入點
     */
    @Pointcut("@annotation(com.example.multidb.anno.DbName)")
    public void chooseDbPointCut(){ }

    @Before("@annotation(dbName)")
    public void changeDataSource(JoinPoint joinPoint, DbName dbName) {
        String dbid = dbName.value();
        if (!DynamicDataSourceContextHolder.isContainsDataSource(dbid)) {
            //joinPoint.getSignature() ：獲取連接點的方法簽名對象
            System.out.println("數據源 " + dbid + " 不存在使用默認的數據源 -> " + joinPoint.getSignature());
        } else {
            System.out.println("使用數據源：" + dbid);
            //向當前線程設置使用的數據源信息。
            //當業務操作時，會由AbstractRoutingDataSource的determineCurrentLookupKey方法，返回從哪個數據源獲取連接。
            //又因DynamicDataSource重寫了determineCurrentLookupKey方法，返回的是ThreadLocal<String>的值
            //所以這樣設置能夠決定哪個數據源起作用
            DynamicDataSourceContextHolder.setDataSourceType(dbid);
        }
    }

    @After("@annotation(dbName)")
    public void clearDataSource(JoinPoint joinPoint, DbName dbName) {
        System.out.println("清除數據源 " + dbName.value() + " ! - start");
        DynamicDataSourceContextHolder.clearDataSourceType();
    }
}

6. 因DynamicDataSourceRegister的注入需要使用@Import，為了方便，我們創建一個注解

@Target({java.lang.annotation.ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Import({DynamicDataSourceRegister.class})
public @interface EnableDynamicDataSource {

}

然后在啟動類上添加注解@EnableDynamicDataSource 。

四、單元測試

1. 創建DAO

接口

public interface UserDao {
    List<Map<String,Object>> listUsers();
}

實現類：

@Repository
public class UserDaoImpl implements UserDao{

    @Autowired
    private JdbcTemplate jdbcTemplate;

    @Override
    public List<Map<String, Object>> listUsers() {
        String sql = "select id, name, age from user";
        return jdbcTemplate.queryForList(sql);

    }
}

2. 創建Service

接口：

public interface UserService {
    List<Map<String,Object>> listUsers();
}

實現類：

@Service
public class UserServiceImpl implements UserService{

    @Autowired
    private UserDao userDao;

    @DbName(value = "tb2") //指定數據源，如果不設置，則默認是default
    @Override
    public List<Map<String, Object>> listUsers() {
        return userDao.listUsers();
    }
}

3. 測試類

@RunWith(SpringRunner.class)
@SpringBootTest
class MultidbApplicationTests {

    @Autowired
    private UserService userService;

    @Test
    void contextLoads() {
        List<Map<String, Object>> userList = userService.listUsers();
        if(null != userList && userList.size()>0){
            for(Map<String,Object> um : userList){
                System.out.println("id:" + um.get("id"));
                System.out.println("name:" + um.get("name"));
                System.out.println("age:" + um.get("age"));
            }
        }
    }
}