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