【Spring Boot】使用 AbstractRoutingDataSource 结合 AOP 实现通过注解动态切换数据源

在项目涉及到多个库的操作时，那么就要设置多个数据源，此时就涉及到数据源的动态切换问题。

本文提供一种通过使用 AbstractRoutingDataSource 结合AOP 实现通过注解动态切换数据源的思路与大致实现，并给出一个Demo。

大致步骤

0. 什么是 AbstractRoutingDataSource
1. application.yaml 添加数据源配置信息
2. 数据源相关配置类
   1. 数据源属性配置类 DruidProperties
   2. 自定义数据源注解 @DataSource
   3. 保存数据源名称工具类 DynamicDataSourceContextHolder
   4. 动态数据源类继承 AbstractRoutingDataSource
3. AOP切面类
4. 源码地址

0. 什么是 AbstractRoutingDataSource

Abstract DataSource implementation that routes getConnection() calls to one of various target DataSources based on a lookup key. The latter is usually (but not necessarily) determined through some thread-bound transaction context.

源码上的注解大致意思是：AbstractRoutingDataSource 是抽象 DataSource 的实现， 其 getConnection() 方法根据 lookupKey 键去调用不同的目标数据源。数据源通常（但不一定）是通过某些线程绑定事务上下文去确定。

也就是说实现了这个spring提供的抽象类，就可以在程序运行的时候，根据自定义的规则动态切换数据源。

重写该类提供的抽象方法determineCurrentLookupKey，提供不同的 lookKey 就可以实现当时使用哪个数据源。

1. application.yaml 添加数据源配置信息

首先需要在application.yaml添加数据源的配置信息，可以按照自己的需要配置多个。

这里只配置了两个数据源的主要信息，一个主数据源和一个从数据源。

# 数据源配置
spring:
  datasource:
    type: com.alibaba.druid.pool.DruidDataSource
    driverClassName: com.mysql.cj.jdbc.Driver
    ds:
      # 主数据源
      master:
        url: jdbc:mysql://localhost:3306/test09?useUnicode=true&characterEncoding=utf8&zeroDateTimeBehavior=convertToNull&useSSL=true&serverTimezone=GMT%2B8
        username: root
        password: ***
      # 从数据源
      slave:
        url: jdbc:mysql://localhost:3306/test08?useUnicode=true&characterEncoding=utf8&zeroDateTimeBehavior=convertToNull&useSSL=true&serverTimezone=GMT%2B8
        username: root
        password: ***

2. 数据源相关配置类

1. 数据源属性配置类 DruidProperties

@ConfigurationProperties(prefix = "spring.datasource")
public class DruidProperties {

    private String type;
    private String driverClassName;
    private Map<String, Map<String, String>> ds;
  
    /**
     * 通过外部构造 DruidDataSource 对象
     */
    public DataSource dataSource(DruidDataSource druidDataSource) {
        return druidDataSource;
    }
  
  	//...getter和setter略
}

我们选择了在外部构造 DataSource对象，还需要一个类加载所有的数据源。

@Component
@EnableConfigurationProperties
public class LoadDataSource {

    @Autowired
    DruidProperties druidProperties;

    /**
     * 用于加载所有的数据源
     */
    public Map<String, DataSource> loadAllDataSource() {
        Map<String, DataSource> map = new HashMap<>();
        Map<String, Map<String, String>> ds = druidProperties.getDs();
        try {
            Set<String> keySet = ds.keySet();
            for (String key : keySet) {
                // 使用 DruidDataSourceFactory 工厂根据给定的参数map创建一个数据源
                map.put(key, druidProperties.dataSource((DruidDataSource) DruidDataSourceFactory.createDataSource(ds.get(key))));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return map;
    }
}

2. 自定义数据源注解 @DataSource

接下来我们需要自定义一个注解@DataSource。该注解的作用是加在Service层的类或方法上，通过 value 属性来指定使用哪个数据源。

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.TYPE, ElementType.METHOD})
public @interface DataSource {
    /**
     * 默认值。如果方法上设置了 @DataSource 注解，但是却没有指定value，那么就默认使用 master 数据源
     */
    String value() default DataSourceType.DEFAULT_MASTER_DS;
}

上面使用了一个 DataSourceType 接口保存常量：

public interface DataSourceType {
    /**
     * 主数据库
     */
    String DEFAULT_MASTER_DS = "master";
    /**
     * 从数据库
     */
    String SLAVE_DS = "slave";
}

3. 保存数据源名称工具类 DynamicDataSourceContextHolder

在开头描述什么是AbstractRoutingDataSource时，在其源码上的有这样一句：数据源通常（但不一定）是通过某些线程绑定事务上下文去确定。

这里的 DynamicDataSourceContextHolder 工具类就是用于存储当前线程中使用的 数据源类型。

@Component
public class DynamicDataSourceContextHolder {

    /**
		 * 用于保存当前线程所使用的数据源名称变量
     */
    public static ThreadLocal<String> CONTEXT_HOLDER = new ThreadLocal<>();

  	/**
  	 * 将给定的dataType数据源类型设置到CONTEXT_HOLDER
  	 */
    public static void setDataSourceType(String dataType) {
        CONTEXT_HOLDER.set(dataType);
    }
  
  	/**
  	 * 从CONTEXT_HOLDER中获取数据源
  	 */
    public static String getDataSourceType() {
        return CONTEXT_HOLDER.get();
    }

    /**
     * ThreadLocal 里面的数据 用完之后一定要记得清空掉，
     * 不然那个线程将来去做其他事情就有可能产生内存泄漏的问题
     */
    public static void clearDataSourceType() {
        CONTEXT_HOLDER.remove();
    }

}

4. 动态数据源类继承 AbstractRoutingDataSource

这里是一个比较关键的点，我们需要设计一个类去继承 AbstractRoutingDataSource，在其中实现我们自定义的获取数据源的逻辑。

@Component
public class DynamicDataSource extends AbstractRoutingDataSource {

    public DynamicDataSource(LoadDataSource loadDataSource) {
        // 从LoadDataSource类中获取加载所有配置好的数据源
        Map<String, DataSource> allDs = loadDataSource.loadAllDataSource();
        // 1.将获取到的数据源封装成map，set为目标数据源
        super.setTargetDataSources(new HashMap<>(allDs));
        // 2.设置默认的数据源
        // 因为将来并不是会在所有的方法上都用 @DataSource 方法，所以对于那些没有 @DataSource注解的方法，需要设置一个默认的数据源
        super.setDefaultTargetDataSource(allDs.get(DataSourceType.DEFAULT_MASTER_DS));
        // 调用父类的方法完成所有属性set后的配置
        super.afterPropertiesSet();
    }

    /**
     * 这个方法是重写了父类的方法，从当前线程中获取使用的数据源类型，以此来获取需要的数据源
     */
    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.getDataSourceType();
    }
}

3. AOP切面类

在完成了数据源配置相关的类，就需要设计一个AOP切面类。

在AOP切面类里面要做的事就是对所有加了@DataSource注解的类或方法，获取其给定要使用的数据源类型名称，保存到 DynamicDataSourceContextHolder 类中的 ThreadLocal 变量中。

这样，就可以实现在需要修改数据源的地方使用注解方式去切换数据源。

@Aspect
@Component
@Order(11)
public class DataSourceAspect {

    /**
     * @annotation(org.andy.dynamic_ds.annotation.DataSource) 表示service层方法上有 @DataSource 注解就将该方法拦截下来
     * @within(org.andy.dynamic_ds.annotation.DataSource 表示service层类上面有 @DataSource 注解，就将类中的方法拦截下来
     */
    @Pointcut("@annotation(org.andy.dynamic_ds.annotation.DataSource) " +
            "|| @within(org.andy.dynamic_ds.annotation.DataSource)")
    public void pc() {

    }

    @Around("pc()")
    public Object around(ProceedingJoinPoint pjp) {
        // 获取方法上面的有效注解
        DataSource dataSource = getDataSource(pjp);
        if (dataSource != null) {
            // 获取注解中数据源的名称
            String value = dataSource.value();
            // set的 ThreadLocal 中
            DynamicDataSourceContextHolder.setDataSourceType(value);
        }
        try {
            return pjp.proceed();
        } catch (Throwable throwable) {
            throwable.printStackTrace();
        } finally {
            // 一定要记得在finally中 清除 ThreadLocal 中使用过的数据变量
            DynamicDataSourceContextHolder.clearDataSourceType();
        }
        return null;
    }

    /**
     * 获取数据源注解@DataSoruce
     * 先在方法上面找，找到直接返回，找不到再从类上面的找。
     */
    private DataSource getDataSource(ProceedingJoinPoint pip) {
        MethodSignature signature = (MethodSignature) pip.getSignature();
        // 查找方法上的注解
        DataSource annotation = AnnotationUtils.findAnnotation(signature.getMethod(), DataSource.class);
        if (annotation != null) {
            // 说明方法上面有 @DataSource 注解，直接返回
            return annotation;
        }
        // 否则去类上面找是否存在 @DataSource 注解
        return AnnotationUtils.findAnnotation(signature.getDeclaringType(), DataSource.class);
    }

4. Demo源码地址

Github地址：https://github.com/RunJavaCode/spring-boot-dynamic-datasource.git

---

以上内容来源于网上学习与整理，手动整理一遍写出来比单纯看一遍视频能记忆更深，理解也更深一点。