一、SpringBoot切换数据源

前言

记得刚工作的，看到别人写了一个切换数据源的，觉得很厉害，当时一直想模仿写一个，但是一直没有行动起来，后面也就逐渐忘记了。下面开始记录下实现过程。

实现原理也非常简单，主要包括了3个步骤：

• 继承 AbstractRoutingDataSource
• 创建数据源
• 创建注解和切面逻辑

下面按照这3个主要步骤来实现

1. 继承 AbstractRoutingDataSource

为什么要继承这个类？

Springboot内置了一个AbstractRoutingDataSource，将所有数据源装入map，然后可以根据不同的key返回不同的数据源。当springboot开始执行连接数据库之前会执行determineCurrentLookupKey()方法，这个方法返回的数据将作为key去map中查找相应的数据源。

public class DynamicDataSource extends AbstractRoutingDataSource {
    @Value("${mini-db-router.jdbc.datasource.default}")
    private String defaultDataSource;

    @Override
    protected Object determineCurrentLookupKey() {
        if (DBContextHolder.getDbkey() == null) {
            return defaultDataSource;
        } else {
            return DBContextHolder.getDbkey();
        }
    }
}

pom文件：

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>com.keller</groupId>
    <artifactId>copy-project</artifactId>
    <version>1.0-SNAPSHOT</version>

    <parent>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-parent</artifactId>
        <version>2.3.5.RELEASE</version>
        <relativePath/> <!-- lookup parent from repository -->
    </parent>

    <dependencies>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-test</artifactId>
            <scope>test</scope>
        </dependency>
        <dependency>
            <groupId>org.mybatis.spring.boot</groupId>
            <artifactId>mybatis-spring-boot-starter</artifactId>
            <version>2.1.4</version>
        </dependency>
        <dependency>
            <groupId>mysql</groupId>
            <artifactId>mysql-connector-java</artifactId>
            <version>5.1.47</version>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>
    </dependencies>


</project>

配置文件 application.yml

server:
  port: 8081

# 多数据源路由配置
mini-db-router:
  jdbc:
    datasource:
      dbCount: 2
      tbCount: 4
      default: db00
      routerKey: uId
      list: db01,db02
      db00:
        driver-class-name: com.mysql.jdbc.Driver
        url: jdbc:mysql://127.0.0.1:3306/lottery?useUnicode=true&characterEncoding=utf8
        username: root
        password: 123456
      db01:
        driver-class-name: com.mysql.jdbc.Driver
        url: jdbc:mysql://127.0.0.1:3306/lottery_01?useUnicode=true&characterEncoding=utf8
        username: root
        password: 123456
      db02:
        driver-class-name: com.mysql.jdbc.Driver
        url: jdbc:mysql://127.0.0.1:3306/lottery_02?useUnicode=true&characterEncoding=utf8
        username: root
        password: 123456

mybatis:
  mapper-locations: classpath:/mybatis/mapper/*.xml
  config-location:  classpath:/mybatis/config/mybatis-config.xml

2. 创建数据源

从配置文件中读取数据库信息创建数据源，通过实现EnvironmentAware接口，获取yml中的配置。

@Configuration
public class DataSourceAutoConfig implements EnvironmentAware {
    /**
     * 数据源配置组
     */
    private Map<String, Map<String, String>> dataSourceMap = new HashMap<>();

    @Value("${mini-db-router.jdbc.datasource.default}")
    private String defaultDataSouceName;

    /**
     * 创建数据源
     */
    @Bean
    public DataSource createDataSouce() {
        Map<Object, Object> targetDataSouce = new HashMap<>();
        for (String dbInfo : dataSourceMap.keySet()) {
            DataSource ds = createDataSource(dataSourceMap.get(dbInfo));
            targetDataSouce.put(dbInfo, ds);
        }

        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        dynamicDataSource.setTargetDataSources(targetDataSouce);
        dynamicDataSource.setDefaultTargetDataSource(createDataSource(dataSourceMap.get(defaultDataSouceName)));

        return dynamicDataSource;
    }

    private DataSource createDataSource(Map<String, String> attributes) {
        try {
            DataSourceProperties dataSourceProperties = new DataSourceProperties();
            dataSourceProperties.setUrl(attributes.get("url").toString());
            dataSourceProperties.setUsername(attributes.get("username").toString());
            dataSourceProperties.setPassword(attributes.get("password").toString());

            String driverClassName = attributes.get("driver-class-name") == null ? "com.zaxxer.hikari.HikariDataSource" : attributes.get("driver-class-name").toString();
            dataSourceProperties.setDriverClassName(driverClassName);

            String typeClassName = attributes.get("type-class-name") == null ? "com.zaxxer.hikari.HikariDataSource" : attributes.get("type-class-name").toString();
            DataSourceBuilder<?> dataSourceBuilder = dataSourceProperties.initializeDataSourceBuilder();
            DataSource ds = dataSourceProperties.initializeDataSourceBuilder().type((Class<DataSource>) Class.forName(typeClassName)).build();
            return ds;
        } catch (ClassNotFoundException e) {
            throw new IllegalArgumentException("can not find datasource type class by class name", e);
        }
    }

    /**
     * 获取数据库配置信息
     * @param environment environment
     */
    @Override
    public void setEnvironment(Environment environment) {
        String prefix = "mini-db-router.jdbc.datasource.";
        String dataSources = environment.getProperty(prefix + "list");
        List<String> dataSourceList = Arrays.stream(dataSources.split(",")).collect(Collectors.toList());
        dataSourceList.add("db00");
        for (String source : dataSourceList) {
            Map<String, String> dbConfig = new HashMap<>();
            String driverClassName = environment.getProperty(prefix + source + ".driver-class-name");
            String url = environment.getProperty(prefix + source + ".url");
            String username = environment.getProperty(prefix + source + ".username");
            String password = environment.getProperty(prefix + source + ".password");
            dbConfig.put("driver-class-name", driverClassName);
            dbConfig.put("url", url);
            dbConfig.put("username", username);
            dbConfig.put("password", password);
            dataSourceMap.put(source, dbConfig);
        }
    }
}

3. 创建注解和切面逻辑

自定义一个注解和切面，在切面内获取到执行的数据库，然后执行。

注解：

@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD, ElementType.TYPE})
public @interface SwitchDS {
    String key() default "";
}

切面：记得在上面加@Component

@Aspect
@Component
public class DBAspect {

    @Pointcut("@annotation(com.keller.lottery.route.annotation.SwitchDS)")
    public void aopPoint(){
    }

    @Around("aopPoint() && @annotation(switchDS)")
    public Object doSwitch(ProceedingJoinPoint jp, SwitchDS switchDS) throws Throwable {
        try {
            // 获取数据库类型，如果注解上有从注解上获取，没有从参数的最后一个获取
            String dbType;
            if (!"".equals(switchDS.key())) {
                dbType = switchDS.key();
            } else {
                Object[] args = jp.getArgs();
                dbType = (String)args[args.length - 1];
            }
            DBContextHolder.setDbkey(dbType);
            return jp.proceed();
        } finally {
            DBContextHolder.clear();
        }
    }
}

4. 测试

@Mapper
public interface RuleTreeDao {
    @SwitchDS
    RuleTree queryRuleTreeByTreeId(Long id, String dbType);
}

@RunWith(SpringRunner.class)
@SpringBootTest
public class SwitchTest {

    @Resource
    RuleTreeDao ruleTreeDao;

    @Test
    public void test1() {
        RuleTree ruleTree = ruleTreeDao.queryRuleTreeByTreeId(2110081903l, "db01");
        System.out.println(ruleTree);
    }
}

二、整合成SpringBoot starter

在实际开发中，很多工程都有切换数据源的需求，为了方便，可以将其弄成SpringBoot starter，其他地方需要引入依赖即可。

项目结构：

1. 编写spring.factories

在 resources下的 META-INF下面新建 spring.factories文件，里面写上要自动配置的类。

org.springframework.boot.autoconfigure.EnableAutoConfiguration=com.keller.route.config.DataSourceAutoConfig

然后maven clean install打包，其他项目需要引入这个依赖即可。

注意点：

需要在自动配置类 DataSourceAutoConfig中将切面的 bean注入，不然引入的时候找不到切面这个bean

    @Bean
    public DBAspect pointBean() {
        return new DBAspect();
    }