Spring Boot 实现数据源的动态切换

一、AbstractRoutingDataSource

AbstractRoutingDataSource 是 Spring Boot 提供的可根据指定路由选择当前的数据源，我们可以在执行之前，设置对应的数据源。它提供了一个抽象的方法 determineCurrentLookupKey()，这个方法决定使用哪个数据源。

public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {
    // 配置的所有数据源
    @Nullable
    private Map<Object, Object> targetDataSources;
    // 默认使用的数据源
    @Nullable
    private Object defaultTargetDataSource;
    private boolean lenientFallback = true;
    private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();
    @Nullable
    private Map<Object, DataSource> resolvedDataSources;
    @Nullable
    private DataSource resolvedDefaultDataSource;


    // 决定使用哪个数据源
    @Nullable
    protected abstract Object determineCurrentLookupKey();
}

二、代码实现

2.1 依赖的 jar 包

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

        <dependency>
            <groupId>mysql</groupId>
            <artifactId>mysql-connector-java</artifactId>
            <scope>runtime</scope>
        </dependency>

        <dependency>
            <groupId>org.mybatis.spring.boot</groupId>
            <artifactId>mybatis-spring-boot-starter</artifactId>
            <version>2.2.0</version>
        </dependency>

        <!-- 通过注解来指定使用的数据源 -->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>

2.2 Spring 的配置文件

这里配置 blog_test 和 blog_test1 两个数据源，通过自定义注解进行动态切换。

spring:
  datasource:
    type: com.alibaba.druid.pool.DruidDataSource
    driver-class-name: com.mysql.cj.jdbc.Driver
    druid:
      first:
        url: jdbc:mysql://127.0.0.1:3306/blog_test?useUnicode=true&characterEncoding=UTF-8&allowMultiQueries=true
        username: root
        password: 123456
      second:
        url: jdbc:mysql://127.0.0.1:3306/blog_test1?useUnicode=true&characterEncoding=UTF-8&allowMultiQueries=true
        username: root
        password: 123456

2.3 自定义注解，通过注解指定使用的数据源

/**
 * 数据源开关，用来指定使用哪个数据源
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface DataSourceSwitch {
    DataSourceNames name() default DataSourceNames.FIRST;
}

2.4 自定义枚举，罗列出可以使用的数据源

/**
 * 数据源名称枚举
 */
public enum DataSourceNames {
    FIRST("first"), SECOND("second");

    private final String name;

    DataSourceNames(String name) {
        this.name = name;
    }
}

2.5 实现AbstractRoutingDataSource

/**
 * 重写determineCurrentLookupKey()，实现选择数据源的方式
 */
public class DynamicDataSource extends AbstractRoutingDataSource {
    // 每个线程指定各自需要使用的数据源名称
    private static final ThreadLocal<DataSourceNames> CONTEXT_HOLDER = new ThreadLocal<>();

    public DynamicDataSource(Map<Object, Object> targetDataSources, Object defaultTargetDataSource) {
        // 设置默认的数据源
        super.setDefaultTargetDataSource(defaultTargetDataSource);
        // 设置配置了的所用数据源
        super.setTargetDataSources(targetDataSources);
        // 将targetDataSources复制给resolvedDataSources;
        // 将defaultTargetDataSource复制给resolvedDefaultDataSource;
        super.afterPropertiesSet();
    }

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

    private void setDataSource(DataSourceNames dataSource) {
        CONTEXT_HOLDER.set(dataSource);
    }

    private DataSourceNames getDataSource() {
        return CONTEXT_HOLDER.get();
    }

    private void clearDataSource() {
        CONTEXT_HOLDER.remove();
    }
}

2.6 通过aop的方式，通过注解修改ThreadLocal的值

/**
 * 处理DataSourceSwitch注解，实现动态切换数据源
 */
@Aspect
@Component
public class DataSourceAspect implements Ordered {

    @Pointcut("@annotation(com.learn.blog.demo.DataSourceSwitch)")
    public void dataSourcePointCut() {};

    @Around("dataSourcePointCut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        MethodSignature methodSignature = (MethodSignature) point.getSignature();
        // 获取方法
        Method method = methodSignature.getMethod();
        // 获取方法上指定的注解
        DataSourceSwitch dataSourceSwitch = method.getAnnotation(DataSourceSwitch.class);
        DynamicDataSource.setDataSource(dataSourceSwitch.name());

        try {
            return point.proceed();
        } finally {
            // 最后清空
            DynamicDataSource.clearDataSource();
        }

    }

    @Override
    public int getOrder() {
        return 1;
    }
}

2.7 多数据源配置类

/**
 * 多数据源配置类
 */
@Configuration
public class DynamicDataSourceConfig {

    /**
     * 实例化数据源first
     * @return DataSource
     */
    @Bean
    @ConfigurationProperties("spring.datasource.druid.first")
    public DataSource firstDataSource() {
        return DruidDataSourceBuilder.create().build();
    }

    /**
     * 实例化数据源second
     * @return DataSource
     */
    @Bean
    @ConfigurationProperties("spring.datasource.druid.second")
    public DataSource secondDataSource() {
        return DruidDataSourceBuilder.create().build();
    }

    /**
     * 实例化DynamicDataSource
     *
     * @param firstDataSource firstDataSource
     * @param secondDataSource secondDataSource
     * @return DynamicDataSource
     */
    @Bean
    @Primary
    public DynamicDataSource dynamicDataSource(DataSource firstDataSource, DataSource secondDataSource) {
        HashMap<Object, Object> targetDataSources = new HashMap<>();
        targetDataSources.put(DataSourceNames.FIRST, firstDataSource);
        targetDataSources.put(DataSourceNames.SECOND, secondDataSource);

        return new DynamicDataSource(targetDataSources, firstDataSource);
    }

    /**
     * 注入SqlSessionFactory

     * @param dynamicDataSource dynamicDataSource
     * @return SqlSessionFactory
     * @throws Exception Exception
     */
    @Bean
    public SqlSessionFactory sqlSessionFactory(DynamicDataSource dynamicDataSource) throws Exception {
        SqlSessionFactoryBean sqlSessionFactoryBean = new SqlSessionFactoryBean();
        sqlSessionFactoryBean.setDataSource(dynamicDataSource);
        return sqlSessionFactoryBean.getObject();
    }

}

三、测试动态切换数据源

-- 在 blog_test 和 blog_test1中分别创建 user 表
create table user(
    id int,
    name varchar(50),
    password varchar(256),
    primary key(id)
);

3.1 Controller层

@RestController
@RequestMapping("/user")
public class UserController {
    @Autowired
    private UserService userService;

    @RequestMapping("/create")
    public String createUser(User user) {
        try {
            userService.createUser(user);
            return "success";
        } catch (Exception e) {
            System.out.println(e);
            return "failed";
        }
    }

}

3.2 Service层

@Service
public class UserService {
    @Autowired
    private UserDao userDao;

    // 默认是 DataSourceNames.FIRST，即 blog_test
    // 这里指定向 DataSourceNames.SECOND写，即 blot_test1
    @DataSourceSwitch(name = DataSourceNames.SECOND)
    public void createUser(User user) {
        userDao.add(user);
    }

}

3.3 Dao层

@Mapper
public interface UserDao {
    @Insert("insert into user values (#{id}, #{name}, #{password})")
    void add(User user);
}

3.4 测试结果