开发环境
- jdk 11.0.10
- SpringBoot 2.6.2
- Idea
本文的代码实现了动态数据源切换的starter工具类。
SpringBoot提供了AbstractRoutingDataSource类用于确定运行时使用哪一个数据源，通过继承该类能够获得动态切换数据源的功能。
首先需要一个切换数据源的上下文辅助类DynamicDataSourceContextHolder,该类中使用了ThreadLocal来为每个请求线程设置需要切换的数据源的key
public final class DynamicDataSourceContextHolder {
/**
* 实际的数据源
*/
private static final Map DS_MAP = new ConcurrentHashMap<>();
/**
* 数据源类型-数据源key列表映射
*/
private static final Map> KEY_MAP = new ConcurrentHashMap<>();
/**
* 默认主数据源的key
*/
private static final ThreadLocal CONTEXT = ThreadLocal.withInitial(() -> null);
/**
* 添加数据源
* @param beanName 数据源BeanName
* @param type 数据源类型
* @param dataSource 数据源
*/
public synchronized static void addMap(String beanName, DataSourceType type, DataSource dataSource) {
var list = KEY_MAP.get(type);
if (list == null) {
list = new ArrayList();
}
list.add(beanName);
KEY_MAP.put(type, list);
// 实际数据源的缓存
DS_MAP.put(beanName, dataSource);
}
/**
* 获取数据源
* @param key 数据源key
* @return 数据源
*/
protected synchronized static DataSource getDataSource(String key) {
return DS_MAP.get(key);
}
/**
* 将第一个从数据源转换为主数据源
* @return 主数据源的key
*/
public synchronized static String convertSlaveToMaster() {
var sList = KEY_MAP.get(DataSourceType.SLAVE);
var mList = new ArrayList(1);
String key = sList.remove(0);
mList.add(key);
KEY_MAP.put(DataSourceType.MASTER, mList);
return key;
}
/**
* 当前数据源的key
* @return 当前数据源的key
*/
public static String current() {
return CONTEXT.get();
}
/**
* 切换为从数据源(随机)
*/
public static void slave() {
var list = KEY_MAP.get(DataSourceType.SLAVE);
if (list == null || list.isEmpty()) {
CONTEXT.set(null);
} else {
// 简单随机
CONTEXT.set(list.get(new Random().nextInt(128) % list.size()));
}
}
/**
* 切换为从数据源(指定)
* @param key 指定数据源key
*/
public static void slave(String key) {
if (KEY_MAP.get(DataSourceType.SLAVE).contains(key)) {
CONTEXT.set(key);
} else {
throw new IllegalArgumentException("no slave datasource matched!...");
}
}
/**
* 切换为主数据源
*/
public static void master() {
// 一定会有一个主数据源
CONTEXT.set(KEY_MAP.get(DataSourceType.MASTER).get(0));
}
/**
* 指定数据源
* @param key 指定数据源key
*/
public static void specify(String key) {
if (DS_MAP.containsKey(key)) {
CONTEXT.set(key);
} else {
throw new IllegalArgumentException("key not matched!");
}
}
/**
* 重置数据源
*/
public static void reset() {
// 如果找不到对应的key,AbstractRoutingDataSource中会使用默认的数据源
// 默认数据源需要在设置AbstractRoutingDataSource时指定
CONTEXT.remove();
}
}
需要注意的是，每次切换数据源后都必须将线程变量中的数据源key清空，否则会影响下一个请求。
这个类中，还持有两个Map, 一个是数据源类型-数据源key列表映射,在外部切换时，直接指定数据源类型即可，这个辅助类会根据类型找到对应所有可用的数据源。另一个是保存实际的数据源的Map,辅助类在实际寻找数据源时利用key找到对应的数据源。具体实现方式见slave()方法。
下面需要实现AbstractRoutingDataSource类
public class DynamicDataSource extends AbstractRoutingDataSource {
private static final Logger logger = LoggerFactory.getLogger(DynamicDataSource.class);
private DataSource defaultTargetDataSource;
@Override
public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
logger.info("default dataSource has been setup...");
this.defaultTargetDataSource = (DataSource) defaultTargetDataSource;
}
/**
* 确定当前需要获取的数据源key-父类中根据key获取对应的数据源
* @return 数据源的key
*/
@Override
protected String determineCurrentLookupKey() {
String key = DynamicDataSourceContextHolder.current();
logger.info("Current DataSource Is {}", key);
return key;
}
/**
* 重写查找数据源的方法-一般重写determineCurrentLookupKey已经足够满足业务需求了
* @return 确定的数据源
*/
@Override
protected DataSource determineTargetDataSource() {
// 由于重写了determineTargetDataSource方法,所以此处determineCurrentLookupKey的作用可有可无
String lookupKey = this.determineCurrentLookupKey();
if (lookupKey == null || DynamicDataSourceContextHolder.getDataSource(lookupKey) == null) {
logger.info("use the default dataSource...");
return this.defaultTargetDataSource;
}
var ds = DynamicDataSourceContextHolder.getDataSource(lookupKey);
// ThreadLocal的key为弱引用,可能已经被GC回收了
if (ds == null) {
// 清除缓存
logger.warn("dataSource key has been garbage collected... current dataSource has been convert to the default dataSource");
DynamicDataSourceContextHolder.reset();
return this.defaultTargetDataSource;
}
return ds;
}
/**
* 重写afterPropertiesSet,让数据源可以在容器初始化后再设置
*/
@Override
public void afterPropertiesSet() {
logger.info("AbstractRoutingDataSource has been override...default setup will not be executed...");
}
正常情况下实现determineCurrentLookupKey方法其实已经够了，实际运行中就是根据这个方法的返回值去找对应的数据源，找不到则使用默认的defaultTargetDataSource。
但是这边考虑到一主多从的情况，所以将determineTargetDataSource也重写了。具体逻辑也不复杂，具体参考贴出的代码。
注意这里还重写了afterPropertiesSet方法，因为正常情况下afterPropertiesSet会在容器生成Bean后进行调用，而父类的afterPropertiesSet方法中会进行一些处理，如下：
@Override
public void afterPropertiesSet() {
if (this.targetDataSources == null) {
throw new IllegalArgumentException("Property 'targetDataSources' is required");
}
this.resolvedDataSources = CollectionUtils.newHashMap(this.targetDataSources.size());
this.targetDataSources.forEach((key, value) -> {
Object lookupKey = resolveSpecifiedLookupKey(key);
DataSource dataSource = resolveSpecifiedDataSource(value);
this.resolvedDataSources.put(lookupKey, dataSource);
});
if (this.defaultTargetDataSource != null) {
this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
}
}
具体可看
AbstractRoutingDataSource源码
由于我们是要做一个工具类，所以需要开放权限给用户自定义，所以初始化时没有设置数据源集合，默认处理会导致异常，这里重写afterPropertiesSet方法避免出错。
数据类型的枚举类如下：
public enum DataSourceType {
/**
* 默认类型,和MASTER等价
*/
DEFAULT("DEFAULT"),
/**
* 主数据源,一般用于写
*/
MASTER("MASTER"),
/**
* 从数据源,一般用于都
*/
SLAVE("SLAVE");
/**
* 数据源类型名称
*/
private String name;
private DataSourceType(String name) {
this.name = name;
}
/**
* 获取类型名称
* @return 类型名称
*/
public String typeName() {
return this.name;
}
}
完成核心功能后，可以开始添加注解的功能。
自定义注解如下：
@Documented
@Retention(RUNTIME)
@Target(value = {ElementType.TYPE, ElementType.METHOD})
public @interface Dynamic {
/**
* 指定数据源key
* @return 数据源key
*/
String target() default "";
/**
* 指定数据源类型(随机获取)
* @return 数据源类型
*/
DataSourceType type() default DataSourceType.DEFAULT;
}
切面类如下：
@Aspect
@Order(-1) // 高优先级-必须要比@Transacation高
public class DynamicSwitchAspect {
private static final Logger logger = LoggerFactory.getLogger(DynamicSwitchAspect.class);
/**
* 切点
*/
@Pointcut("@annotation(cn.t.dynamic.switcher.annotation.Dynamic) || @within(cn.t.dynamic.switcher.annotation.Dynamic)")
public void pointCut() {}
/**
* 方法执行前切换数据源
* @param joinPoint 连接点信息
*/
@Before("pointCut()")
public void before(JoinPoint joinPoint) {
// 此处注解信息一定存在于类上或者方法上
// 优先使用方法上的注解
Dynamic d = ((MethodSignature)joinPoint.getSignature()).getMethod().getAnnotation(Dynamic.class);
if (d == null) {
d = joinPoint.getTarget().getClass().getAnnotation(Dynamic.class);
}
String target = d.target();
DataSourceType type = d.type();
if (target.isEmpty() && type.equals(DataSourceType.DEFAULT)) {
// 使用默认的数据源
DynamicDataSourceContextHolder.master();
} else if (!target.isEmpty() && type.equals(DataSourceType.DEFAULT)) {
// 指定数据源
DynamicDataSourceContextHolder.specify(target);
} else if (target.isEmpty() && type.equals(DataSourceType.SLAVE)) {
// 随机选取从数据源
DynamicDataSourceContextHolder.slave();
} else if (!target.isEmpty() && type.equals(DataSourceType.SLAVE)) {
// 指定从数据源
DynamicDataSourceContextHolder.slave(target);
} else {
// 排除所有情况后
// 主数据源
DynamicDataSourceContextHolder.master();
}
logger.info("datasource conversion completed");
}
/**
* 方法执行后清空切换key
* @param joinPoint 连接点信息
*/
@After("pointCut()")
public void after(JoinPoint joinPoint) {
logger.info("datasource key will be reset...");
DynamicDataSourceContextHolder.reset();
}
}
由于注解支持类和方法，所以切面类中需要仔细判断。
由于容器启动过程中不会指定数据源集合，所以需要在容器启动后进行增强处理，在增强处理中配置数据源集合：
public class DynamicContextAware implements ApplicationContextAware {
private static final Logger logger = LoggerFactory.getLogger(DynamicContextAware.class);
/**
* 容器后处理-在最后给动态数据源设置可用的数据源
* @param applicationContext Spring容器上下文
* @throws BeansException 异常
*/
@Override
public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
// 在所有Bean初始化完成后获取动态数据源,最后进行设置,保证自定义的数据源已经被初始化完成了
DynamicDataSource ds = applicationContext.getBean(DynamicSwitcherConfiguration.DYNAMIC_SOURCE_BEAN_NAME, DynamicDataSource.class);
// 获取所有的数据源实例(BeanName-DataSource)(需要去除动态数据源-动态数据源也是DataSource类型)
Map beans = applicationContext.getBeansOfType(DataSource.class);
// 去除自身
beans.remove(DynamicSwitcherConfiguration.DYNAMIC_SOURCE_BEAN_NAME);
if (beans.isEmpty()) {
logger.warn("it seems to you have not setup the dataSource...");
return;
}
// 配置多数据源Map
Map dsMap = CollectionUtils.newHashMap(beans.size());
var hasPrimary = false;
for (Map.Entry entry : beans.entrySet()) {
var v = entry.getValue();
var k = entry.getKey();
// DataSource是在配置类中使用@Bean方法生成的,DataSource类上并没有Primary注解,需要使用Spring容器上下文来获取
// var p = v.getClass().getAnnotation(Primary.class);
if (this.isPrimary(applicationContext, k)) {
if (!hasPrimary) {
// 设置主数据源
ds.setDefaultTargetDataSource(v);
hasPrimary = true;
DynamicDataSourceContextHolder.addMap(k, DataSourceType.MASTER, v);
} else {
logger.info("already has a master...{} will become a slave...", k);
DynamicDataSourceContextHolder.addMap(k, DataSourceType.SLAVE, v);
}
} else {
DynamicDataSourceContextHolder.addMap(k, DataSourceType.SLAVE, v);
}
dsMap.put(k, v);
}
// 未指定默认数据源则使用第一个
if (!hasPrimary) {
String key = DynamicDataSourceContextHolder.convertSlaveToMaster();
logger.warn("do not find a primary dataSource...the first dataSource whose name is {} will become the master...", key);
ds.setDefaultTargetDataSource(beans.get(key));
}
// 设置多数据源Map
// ds.setTargetDataSources(dsMap);
}
/**
* 判断当前数据源是否被Primary注解标注(由于DataSource是在配置类中使用@Bean方法生成的,普通的getClass()再获取注解将不适用)
* 而使用容器的context的方法可以获取到
* @param applicationContext 容器上下文
* @param key 比对的数据源key
* @return 是否标注了primary注解
*/
private boolean isPrimary(ApplicationContext applicationContext, String key) {
var beanMap = applicationContext.getBeansWithAnnotation(Primary.class);
return beanMap.containsKey(key);
}
}
最后，配置自动配置类：
@Configuration
public class DynamicSwitcherConfiguration {
public static final String DYNAMIC_SOURCE_BEAN_NAME = "dynamicDataSource";
/**
* 创建切面Bean
* @return
*/
@Bean
public DynamicSwitchAspect dynamicSwitchAspect() {
return new DynamicSwitchAspect();
}
/**
* 创建后处理Bean
* @return
*/
@Bean
public DynamicContextAware dynamicContextAware() {
return new DynamicContextAware();
}
/**
* 动态数据源上下文-内部应该包含一个数据源key-数据源对象的键值对Map
* mybatis中直接注入这个bean即可
* @return
*/
@Bean(name = DYNAMIC_SOURCE_BEAN_NAME)
public DynamicDataSource dynamicDataSource() {
return new DynamicDataSource();
}
}
使用方式
打包后，引入需要该功能的模块中
cn.t.dynamic
dataSource-switcher-spring-boot-starter
1.0.0
在自己模块中注入动态数据源
@Autowired
private DynamicDataSource dynamicDataSource;
在为类似mybatis框架配置数据源时,使用dynamicDataSource即可
多数据源配置可参考官网
https://docs.spring.io/spring-boot/docs/2.4.2/reference/html/howto.html#howto-two-datasources
-
数据源需要自己配置生成，配置过程中可高度定制化。
数据源需要自己配置生成，配置过程中可高度定制化。
//ds1...
@Bean("prop1")
@ConfigurationProperties(prefix = "spring.datasource.datasource1")
public DataSourceProperties dsProp1() {
return new DataSourceProperties();
}
@Bean(name = "dataSource1")
@Primary // 该注解标注的数据源会被认为是master
public DataSource dataSource1() {
// 可以自定义一些属性,比如连接池的设置
return dsProp1().initializeDataSourceBuilder().build();
}
// ds2....
配置事务和sessionFactory时注入动态数据源
// 以事务配置举例
@Bean(name = "sf1")
@Primary
public PlatformTransactionManager txManager1(DynamicDataSource dataSource) {
DataSourceTransactionManager transactionManager = new DataSourceTransactionManager(dataSource);
// 事务10秒超时
transactionManager.setDefaultTimeout(10);
return transactionManager;
}
切换数据源
切换数据源(编码方式)
-
切换
masterDynamicDataSourceContextHolder.master() -
切换
slaveDynamicDataSourceContextHolder.slave() -
指定
slaveDynamicDataSourceContextHolder.slave(String) -
指定存在的数据源
DynamicDataSourceContextHolder.specify(String)
public voud func() {
// 切换数据源
DynamicDataSourceContextHolder.master();
// DynamicDataSourceContextHolder.slave();
// do something
// ...
// !!!清除缓存
DynamicDataSourceContextHolder.reset();
}
完成逻辑后一定要调用DynamicDataSourceContextHolder.reset()清除缓存
切换数据源(注解方式)
-
切换
master@Dynamic public void demo() { }@Dynamic(type = DataSourceType.MASTER) public void demo() { }// 当类型指定为MASTER时,key任意,因为主数据源只有一个 @Dynamic(target = "any key", type = DataSourceType.MASTER) public void demo() { } -
切换
slave// 指定一个从数据源,target必须存在且类型为从数据源 @Dynamic(target = "ds2", type = DataSourceType.SLAVE) public void demo() { }// 随机选取一个从数据源 @Dynamic(type = DataSourceType.SLAVE) public void demo() { } -
指定数据源
// 指定一个数据源,target必须存在 @Dynamic(target = "ds2") public void demo() { }
注意
DataSourceType.DEFAULT不需要手动设置