SpringBoot 自定义动态数据源

1. 原理

动态数据源，本质上是把多个数据源存储在一个 Map 中，当需要使用某一个数据源时，使用 key 获取指定数据源进行处理。而在 Spring 中已提供了抽象类 AbstractRoutingDataSource 来实现此功能，继承 AbstractRoutingDataSource 类并覆写其 determineCurrentLookupKey() 方法监听获取 key 即可，该方法只需要返回数据源 key 即可，也就是存放数据源的 Map 的 key。

因此，我们在实现动态数据源的，只需要继承它，实现自己的获取数据源逻辑即可。AbstractRoutingDataSource 顶级继承了 DataSource，所以它也是可以做为数据源对象，因此项目中使用它作为主数据源。

1.1. AbstractRoutingDataSource 源码解析

        public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {
        // 目标数据源 map 集合，存储将要切换的多数据源 bean 信息，可以通过 setTargetDataSource(Map<Object, Object> mp) 设置
        @Nullable
        private Map<Object, Object> targetDataSources;
        // 未指定数据源时的默认数据源对象，可以通过 setDefaultTargetDataSouce(Object obj) 设置
        @Nullable
        private Object defaultTargetDataSource;
  ...
        // 数据源查找接口，通过该接口的 getDataSource(String dataSourceName) 获取数据源信息
        private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();
        //解析 targetDataSources 之后的 DataSource 的 map 集合
        @Nullable
        private Map<Object, DataSource> resolvedDataSources;
        @Nullable
        private DataSource resolvedDefaultDataSource;

        //将 targetDataSources 的内容转化一下放到 resolvedDataSources 中，将 defaultTargetDataSource 转为 DataSource 赋值给 resolvedDefaultDataSource
        public void afterPropertiesSet() {
            //如果目标数据源为空，会抛出异常，在系统配置时应至少传入一个数据源
            if (this.targetDataSources == null) {
                throw new IllegalArgumentException("Property 'targetDataSources' is required");
            } else {
                //初始化 resolvedDataSources 的大小
                this.resolvedDataSources = CollectionUtils.newHashMap(this.targetDataSources.size());
                //遍历目标数据源信息 map 集合，对其中的 key，value 进行解析
                this.targetDataSources.forEach((key, value) -> {
                    // resolveSpecifiedLookupKey 方法没有做任何处理，只是将 key 继续返回
                    Object lookupKey = this.resolveSpecifiedLookupKey(key);
                    // 将目标数据源 map 集合中的 value 值（Druid 数据源信息）转为 DataSource 类型
                    DataSource dataSource = this.resolveSpecifiedDataSource(value);
                    // 将解析之后的 key，value 放入 resolvedDataSources 集合中
                    this.resolvedDataSources.put(lookupKey, dataSource);
                });
                if (this.defaultTargetDataSource != null) {
                    // 将默认目标数据源信息解析并赋值给 resolvedDefaultDataSource
                    this.resolvedDefaultDataSource = this.resolveSpecifiedDataSource(this.defaultTargetDataSource);
                }

            }
        }

        protected Object resolveSpecifiedLookupKey(Object lookupKey) {
            return lookupKey;
        }

        protected DataSource resolveSpecifiedDataSource(Object dataSource) throws IllegalArgumentException {
            if (dataSource instanceof DataSource) {
                return (DataSource)dataSource;
            } else if (dataSource instanceof String) {
                return this.dataSourceLookup.getDataSource((String)dataSource);
            } else {
                throw new IllegalArgumentException("Illegal data source value - only [javax.sql.DataSource] and String supported: " + dataSource);
            }
        }

        // 因为 AbstractRoutingDataSource 继承 AbstractDataSource，而 AbstractDataSource 实现了 DataSource 接口，所有存在获取数据源连接的方法
        public Connection getConnection() throws SQLException {
            return this.determineTargetDataSource().getConnection();
        }

        public Connection getConnection(String username, String password) throws SQLException {
            return this.determineTargetDataSource().getConnection(username, password);
        }

  // 最重要的一个方法，也是 DynamicDataSource 需要实现的方法
        protected DataSource determineTargetDataSource() {
            Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
            // 调用实现类中重写的 determineCurrentLookupKey 方法拿到当前线程要使用的数据源的名称
            Object lookupKey = this.determineCurrentLookupKey();
            // 去解析之后的数据源信息集合中查询该数据源是否存在，如果没有拿到则使用默认数据源 resolvedDefaultDataSource
            DataSource dataSource = (DataSource)this.resolvedDataSources.get(lookupKey);
            if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
                dataSource = this.resolvedDefaultDataSource;
            }

            if (dataSource == null) {
                throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
            } else {
                return dataSource;
            }
        }

        @Nullable
        protected abstract Object determineCurrentLookupKey();
    }

1.2. 关键类说明

忽略掉 controller/service/entity/mapper/xml介绍。

• application.yml：数据源配置文件。但是如果数据源比较多的话，根据实际使用，最佳的配置方式还是独立配置比较好。
• DynamicDataSourceRegister：动态数据源注册配置文件
• DynamicDataSource：动态数据源配置类，继承自 AbstractRoutingDataSource
• TargetDataSource：动态数据源注解，切换当前线程的数据源
• DynamicDataSourceAspect：动态数据源设置切面，环绕通知，切换当前线程数据源，方法注解优先
• DynamicDataSourceContextHolder：动态数据源上下文管理器，保存当前数据源的 key，默认数据源名，所有数据源 key

1.3. 开发流程

1. 添加配置文件，设置默认数据源配置，和其他数据源配置
2. 编写 DynamicDataSource 类，继承 AbstractRoutingDataSource 类，并实现 determineCurrentLookupKey() 方法
3. 编写 DynamicDataSourceHolder 上下文管理类，管理当前线程的使用的数据源，及所有数据源的 key；
4. 编写 DynamicDataSourceRegister 类通过读取配置文件动态注册多数据源，并在启动类上导入（@Import）该类
5. 自定义数据源切换注解 TargetDataSource，并实现相应的切面，环绕通知切换当前线程数据源，注解优先级（DynamicDataSourceHolder.setDynamicDataSourceKey() > Method > Class）

2. 实现

2.1. 引入 Maven 依赖

<!-- web 模块依赖 -->
<dependency>
 <groupId>org.springframework.boot</groupId>
 <artifactId>spring-boot-starter-web</artifactId>
</dependency>
<!-- spring 核心 aop 模块依赖 -->
<dependency>
 <groupId>org.springframework.boot</groupId>
 <artifactId>spring-boot-starter-aop</artifactId>
</dependency>
<!-- Druid 数据源连接池依赖 -->
<dependency>
 <groupId>com.alibaba</groupId>
 <artifactId>druid-spring-boot-starter</artifactId>
 <version>1.2.8</version>
</dependency>
<!-- mybatis 依赖 -->
<dependency>
 <groupId>org.mybatis.spring.boot</groupId>
 <artifactId>mybatis-spring-boot-starter</artifactId>
 <version>2.2.2</version>
</dependency>
<!-- mysql驱动 -->
<dependency>
 <groupId>mysql</groupId>
 <artifactId>mysql-connector-java</artifactId>
 <version>8.0.24</version>
</dependency>
<!-- lombok 模块依赖 -->
<dependency>
 <groupId>org.projectlombok</groupId>
 <artifactId>lombok</artifactId>
 <optional>true</optional>
</dependency>
<dependency>
 <groupId>org.apache.commons</groupId>
 <artifactId>commons-text</artifactId>
 <version>1.10.0</version>
</dependency>
<dependency>
 <groupId>org.springframework.boot</groupId>
 <artifactId>spring-boot-starter-test</artifactId>
 <scope>test</scope>
</dependency>

2.2. application.yml 配置文件

spring:
  datasource:
 type: com.alibaba.druid.pool.DruidDataSource
    url: jdbc:mysql://localhost:3306/test?useUnicode=true&characterEncoding-utf8&allowPublicKeyRetrieval=true&useSSL=false&serverTimezone=Asia/Shanghai&allowMultiQueries=true
    driver-class-name: com.mysql.cj.jdbc.Driver
    username: root
    password: root
custom:
  datasource:
    names: ds1,ds2
    ds1:
      type: com.alibaba.druid.pool.DruidDataSource
      driver-class-name: com.mysql.cj.jdbc.Driver
      url: jdbc:mysql://localhost:3306/content_center?useUnicode
      username: root
      password: root
    ds2:
   type: com.alibaba.druid.pool.DruidDataSource
      driver-class-name: com.mysql.cj.jdbc.Driver
      url: jdbc:mysql://localhost:3306/trade?useUnicode
      username: root
      password: root

2.3. 创建 DynamicDataSource 继承 AbstractRoutingDataSource 类

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.AllArgsConstructor;

/**
 * @Description: 继承Spring AbstractRoutingDataSource 实现路由切换
 */
@Data
@NoArgsConstructor
@AllArgsConstructor
public class DynamicDataSource extends AbstractRoutingDataSource {

 /**
  * 决定当前线程使用哪种数据源
  * @return 数据源 key
  */
    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.getDataSourceType();
    }
}

2.4. 编写 DynamicDataSourceHolder 类，管理 DynamicDataSource 上下文

import java.util.ArrayList;
import java.util.List;

/**
 * @Description: 动态数据源上下文管理
 */
public class DynamicDataSourceHolder {
    // 存放当前线程使用的数据源类型信息
    private static final ThreadLocal<String> DYNAMIC_DATASOURCE_KEY = new ThreadLocal<String>();
    // 存放数据源 key
    private static final List<String> DATASOURCE_KEYS = new ArrayList<String>();
    // 默认数据源 key
    public static final String DEFAULT_DATESOURCE_KEY = "master";

    //设置数据源
    public static void setDynamicDataSourceType(String key) {
        DYNAMIC_DATASOURCE_KEY.set(key);
    }

    //获取数据源
    public static String getDynamicDataSourceType() {
        return DYNAMIC_DATASOURCE_KEY.get();
    }

    //清除数据源
    public static void removeDynamicDataSourceType() {
        DYNAMIC_DATASOURCE_KEY.remove();
    }

 public static void addDataSourceKey(String key) {
  DATASOURCE_KEYS.add(key)
 }

    /**
     * 判断指定 key 当前是否存在
     *
     * @param key
     * @return boolean
     */
    public static boolean containsDataSource(String key){
        return DATASOURCE_KEYS.contains(key);
    }
}

2.5. 编写 DynamicDataSourceRegister 读取配置文件注册多数据源

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.MutablePropertyValues;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.GenericBeanDefinition;
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.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import org.apache.commons.lang3.StringUtils;
import java.util.Objects;

/**
 * @Description: 注册动态数据源
 * 初始化数据源和提供了执行动态切换数据源的工具类
 * EnvironmentAware（获取配置文件配置的属性值）
 */
public class DynamicDataSourceRegister implements ImportBeanDefinitionRegistrar, EnvironmentAware {
 private static final Logger LOGGER = LoggerFactory.getLogger(DynamicDataSourceRegister.class);
    // 指定默认数据源类型 (springboot2.0 默认数据源是 hikari 如何想使用其他数据源可以自己配置)
    // private static final String DATASOURCE_TYPE_DEFAULT = "com.zaxxer.hikari.HikariDataSource";
    private static final String DEFAULT_DATASOURCE_TYPE = "com.alibaba.druid.pool.DruidDataSource";
    // 默认数据源
    private DataSource defaultDataSource;
    // 用户自定义数据源
    private Map<String, DataSource> customDataSources  = new HashMap<>();

    /**
     * 加载多数据源配置
     * @param env 当前环境
     */
    @Override
    public void setEnvironment(Environment env) {
        initDefaultDataSource(env);
        initCustomDataSources(env);
    }



    /**
     * 初始化主数据源
     * @param env
     */
    private void initDefaultDataSource(Environment env) {
        // 读取主数据源
        Map<String, Object> dsMap = new HashMap<>();
        dsMap.put("type", env.getProperty("spring.datasource.type", DEFAULT_DATASOURCE_TYPE));
        dsMap.put("driver", env.getProperty("spring.datasource.driver-class-name"));
        dsMap.put("url", env.getProperty("spring.datasource.url"));
        dsMap.put("username", env.getProperty("spring.datasource.username"));
        dsMap.put("password", env.getProperty("spring.datasource.password"));
        defaultDataSource = buildDataSource(dsMap);
    }


    /**
     * 初始化更多数据源
     * @param env
     */
    private void initCustomDataSources(Environment env) {
        // 读取配置文件获取更多数据源
        String dsPrefixs = env.getProperty("custom.datasource.names");
        if (!StringUtils.isBlank(dsPrefixs)) {
         for (String dsPrefix : dsPrefixs.split(",")) {
          dsPrefix = fsPrefix.trim()
          if (!StringUtils.isBlank(dsPrefix)) {
           Map<String, Object> dsMap = new HashMap<>();
           dsMap.put("type", env.getProperty("custom.datasource." + dsPrefix + ".type", DEFAULT_DATASOURCE_TYPE));
              dsMap.put("driver", env.getProperty("custom.datasource." + dsPrefix + ".driver-class-name"));
              dsMap.put("url", env.getProperty("custom.datasource." + dsPrefix + ".url"));
              dsMap.put("username", env.getProperty("custom.datasource." + dsPrefix + ".username"));
              dsMap.put("password", env.getProperty("custom.datasource." + dsPrefix + ".password"));
              DataSource ds = buildDataSource(dsMap);
              customDataSources.put(dsPrefix, ds);
          }
         }
        }
    }

    @Override
    public void registerBeanDefinitions(AnnotationMetadata annotationMetadata, BeanDefinitionRegistry registry) {
        Map<Object, Object> targetDataSources = new HashMap<Object, Object>();
        // 将主数据源添加到更多数据源中
        targetDataSources.put(DynamicDataSourceHolder.DEFAULT_DATASOURCE_KEY, defaultDataSource);
        DynamicDataSourceHolder.addDataSourceKey(DynamicDataSourceHolder.DEFAULT_DATASOURCE_KEY);
        // 添加更多数据源
        targetDataSources.putAll(customDataSources);
        for (String key : customDataSources.keySet()) {
            DynamicDataSourceContextHolder.addDataSourceKey(key);
        }

        // 创建 DynamicDataSource
        GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
        beanDefinition.setBeanClass(DynamicDataSource.class);
        beanDefinition.setSynthetic(true);
        MutablePropertyValues mpv = beanDefinition.getPropertyValues();
        mpv.addPropertyValue("defaultTargetDataSource", defaultDataSource);
        mpv.addPropertyValue("targetDataSources", targetDataSources);
        registry.registerBeanDefinition("dataSource", beanDefinition); // 注册到 Spring 容器中
        LOGGER.info("Dynamic DataSource Registry");
    }

    /**
     * 创建 DataSource
     * @param dsMap 数据库配置参数
     * @return DataSource
     */
    public DataSource buildDataSource(Map<String, Object> dsMap) {
        try {
            Object type = dsMap.get("type");
            if (type == null)
                type = DEFAULT_DATASOURCE_TYPE;// 默认DataSource

            Class<? extends DataSource> dataSourceType = (Class<? extends DataSource>)Class.forName((String)type);
            String driverClassName = String.valueOf(dsMap.get("driver"));
            String url = String.valueOf(dsMap.get("url"));
            String username = String.valueOf(dsMap.get("username"));
            String password = String.valueOf(dsMap.get("password"));

            // 自定义 DataSource 配置
            DataSourceBuilder<? extends DataSource> factory = DataSourceBuilder.create()
                    .driverClassName(driverClassName)
                    .url(url)
                    .username(username)
                    .password(password)
                    .type(dataSourceType);
            return factory.build();
        }catch (ClassNotFoundException e) {
            e.printStackTrace();
        }
    }
}

2.6. 在启动器类上添加 @Import，导入 register 类

// 注册动态多数据源
@Import({ DynamicDataSourceRegister.class })
@SpringBootApplication
public class Application {
    public static void main(String[] args) {
        SpringApplication.run(Application.class, args);
    }
}

2.7. 自定义注解 @TargetDataSource

/**
 * 自定义多数据源切换注解
 * 优先级：DynamicDataSourceHolder.setDynamicDataSourceKey() > Method > Class
 */
@Target({ ElementType.METHOD, ElementType.TYPE })
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Inherited
public @interface DataSource
{
    /**
     * 切换数据源名称
     */
    public String value() default DynamicDataSourceHolder.DEFAULT_DATESOURCE_KEY;
}

2.8. 定义切面拦截 @TargetDataSource

import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.Objects;

@Aspect
// 保证在 @Transactional 等注解前面执行
@Order(-1)
@Component
public class DataSourceAspect {

    // 设置 DataSource 注解的切点表达式
    @Pointcut("@annotation(com.ayi.config.datasource.DynamicDataSource)")
    public void dynamicDataSourcePointCut(){

    }

    //环绕通知
    @Around("dynamicDataSourcePointCut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable{
        String key = getDefineAnnotation(joinPoint).value();
        if (!DynamicDataSourceHolder.containsDataSource(key)) {
         LOGGER.error("数据源[{}]不存在，使用默认数据源[{}]", key, DynamicDataSourceHolder.DEFAULT_DATESOURCE_KEY)
         key = DynamicDataSourceHolder.DEFAULT_DATESOURCE_KEY;
        }
        DynamicDataSourceHolder.setDynamicDataSourceKey(key);
        try {
            return joinPoint.proceed();
        } finally {
            DynamicDataSourceHolder.removeDynamicDataSourceKey();
        }
    }

    /**
     * 先判断方法的注解，后判断类的注解，以方法的注解为准
     * @param joinPoint 切点
     * @return TargetDataSource
     */
    private TargetDataSource getDefineAnnotation(ProceedingJoinPoint joinPoint){
        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
        TargetDataSource dataSourceAnnotation = methodSignature.getMethod().getAnnotation(TargetDataSource.class);
        if (Objects.nonNull(methodSignature)) {
            return dataSourceAnnotation;
        } else {
            Class<?> dsClass = joinPoint.getTarget().getClass();
            return dsClass.getAnnotation(TargetDataSource.class);
        }
    }

}