前言

Spring提供了AbstractRoutingDataSource类以方便开发者实现多数据源,看一下AbstractRoutingDataSource#getConnection()的源码:

1
2
3
4
@Override
public Connection getConnection() throws SQLException {
return determineTargetDataSource().getConnection();
}

可以看到在getConnection()方法中是通过调用determineTargetDataSource().getConnection();获取一个连接,继续追踪到determineTargetDataSource()方法:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
protected DataSource determineTargetDataSource() {
Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
//获取key,模板方法模式
Object lookupKey = determineCurrentLookupKey();
//根据key获取DataSource
DataSource dataSource = this.resolvedDataSources.get(lookupKey);
//如果dataSource为空并且启用获取默认dataSource或lookupKey为空时取默认的DataSource
//(lenientFallback默认为true)
if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
dataSource = this.resolvedDefaultDataSource;
}
if (dataSource == null) {//如果dataSource为空抛出异常
throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
}
return dataSource;
}

可以看到determineTargetDataSource()方法先调用determineCurrentLookupKey()获取lookupKey,再通过this.resolvedDataSources.get(lookupKey);取出DataSource,我们看一下determineCurrentLookupKey()方法,发现这个方法是个abstract方法,也就是说子类必须要实现,再看下this.resolvedDataSources是个什么东西:

1
private Map<Object, DataSource> resolvedDataSources;

所以说AbstractRoutingDataSource中维护了一个Map,在getConnection的时候先获取key再从resolvedDataSources中get一个DataSource返回(如果为空时判断是否启用获取默认dataSource再决定是否用默认的dataSource)而获取key的方法determineCurrentLookupKey()由子类实现(模板方法模式)。
所以实现动态多数据源的思路就十分明确了

  1. 定义一个DynamicDataSourceContextHolder用于保存当前线程需要使用的dataSource对应的key
  2. 重写determineCurrentLookupKey()方法将这个key返回

实际上完成上述两个步骤其实就可以实现多数据源了,只要在getConnection前调用DynamicDataSourceContextHolder#setKey方法设置需要使用的dataSource对应的key就可以了。
但通常来说,我们并不希望设置使用那个数据库的代码侵入到我们的业务代码中,所以我们可以利用aop实现:定义一个注解@DataSource和注解切面DataSourceAspect,然后就可以在需要切换数据库的方法上使用注解进行设置。

代码

DataSourceConfiguration

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
/**
* @author lolicom
*/
@Slf4j
@Configuration
public class DataSourceConfiguration implements ApplicationContextAware {

private ApplicationContext applicationContext;

@Override
public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
this.applicationContext = applicationContext;
}

/* configuration */

@Bean
@Qualifier("defaultDataSource")
@ConfigurationProperties(prefix = "spring.datasource.first")
public DataSource dataSource1() {
return DataSourceBuilder.create().build();
}

@Bean
@ConfigurationProperties(prefix = "spring.datasource.second")
public DataSource dataSource2() {
return DataSourceBuilder.create().build();
}

@Bean
@Primary
public DataSource dataSource(Map<String, DataSource> dataSourceMap, @Qualifier("defaultDataSource") DataSource defaultDataSource) {
AbstractRoutingDataSource dataSource = new AbstractRoutingDataSource() {
@Override
protected Object determineCurrentLookupKey() {
return DynamicDataSourceContextHolder.getKey();
}

@Override
public void setTargetDataSources(Map<Object, Object> targetDataSources) {
super.setTargetDataSources(targetDataSources);
DynamicDataSourceContextHolder.setDataSourceMap(targetDataSources);
}

@Override
public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
if (defaultTargetDataSource instanceof String) {
super.setDefaultTargetDataSource(dataSourceMap.get(defaultTargetDataSource));
DynamicDataSourceContextHolder.setDefaultKey((String) defaultTargetDataSource);
} else if (defaultTargetDataSource instanceof DataSource) {
super.setDefaultTargetDataSource(defaultTargetDataSource);
DynamicDataSourceContextHolder.setDefaultKey(resolveSpecifiedLookupKey((DataSource) defaultTargetDataSource));
} else {
log.info("Why am i here?");
}
}

private String resolveSpecifiedLookupKey(DataSource defaultTargetDataSource) {
String[] beanDefinitionNames = applicationContext.getBeanNamesForType(defaultTargetDataSource.getClass());
for (String beanDefinitionName : beanDefinitionNames) {
if (applicationContext.getBean(beanDefinitionName) == defaultTargetDataSource) {
return beanDefinitionName;
}
}
return null;
}
};
dataSource.setTargetDataSources(Collections.unmodifiableMap(dataSourceMap));
dataSource.setDefaultTargetDataSource(defaultDataSource);
return dataSource;
}

}

DynamicDataSourceContextHolder

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
@Slf4j
public class DynamicDataSourceContextHolder {
private final static ThreadLocal<String> KEY = new ThreadLocal<>();
private static Map<Object, Object> targetDataSourceMap;
private static String defaultKey;

public static void setDataSourceMap(Map<Object, Object> targetDataSourceMap) {
DynamicDataSourceContextHolder.targetDataSourceMap = targetDataSourceMap;
}

public static String getKey() {
return Optional.ofNullable(DynamicDataSourceContextHolder.KEY.get())
.orElseGet(() -> DynamicDataSourceContextHolder.defaultKey);
}

public static void setKey(String key) {
DynamicDataSourceContextHolder.KEY.set(targetDataSourceMap.containsKey(key) ? key : DynamicDataSourceContextHolder.defaultKey);
}

public static void remove() {
DynamicDataSourceContextHolder.KEY.remove();
}

public static void setDefaultKey(String defaultKey) {
DynamicDataSourceContextHolder.defaultKey = defaultKey;
log.debug("设置defaultKey:[{}]", defaultKey);
}
}

DataSource注解

1
2
3
4
5
6
7
8
9
10
/**
* @author lolicom
*/
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface DataSource {
String value() default "";
}

DataSourceAspect切面

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
@Slf4j
@Component
@Aspect
@Order(Ordered.LOWEST_PRECEDENCE - 1)
public class DataSourceAspect {

@Pointcut(value = "@annotation(dataSource)", argNames = "dataSource")
public void pointcut(DataSource dataSource) {
}

@Before(value = "pointcut(dataSource)", argNames = "dataSource")
public void before(DataSource dataSource) {
String value = dataSource.value();
DynamicDataSourceContextHolder.setKey(value);
log.debug("使用数据库{}", value);
}

@After("@annotation(cn.griouges.learnspringboot.common.annotation.DataSource)")
public void after() {
DynamicDataSourceContextHolder.remove();
}
}

测试发现如果切面优先级为Ordered.LOWEST_PRECEDENCE时,每次都会在getConnection之后再拦截进行设置key,不符合我们的需求,添加@Order(Ordered.LOWEST_PRECEDENCE - 1)设置切面优先级,在获取数据库连接前进行设置数据库key。

配置完毕

后续还有数据源,只要注册bean到容器中就可以自动添加到AbstractRoutingDataSourcetargetDataSources中,key为bean的name。


测试

application.properties中添加配置

1
2
3
4
5
6
7
8
spring.datasource.first.jdbc-url=jdbc:mysql://localhost:3306/mysite?useUnicode=true&characterEncoding=utf-8
spring.datasource.first.username=root
spring.datasource.first.password=leisure.
spring.datasource.first.driver-class-name=com.mysql.cj.jdbc.Driver
spring.datasource.second.jdbc-url=jdbc:mysql://localhost:3306/test?useUnicode=true&characterEncoding=utf-8
spring.datasource.second.username=root
spring.datasource.second.password=leisure.
spring.datasource.second.driver-class-name=com.mysql.cj.jdbc.Driver

service层的方法上添加注解

1
2
3
4
@DataSource("dataSource2")
public User findUserForLogin(String username, String password) {
return userRepository.findByUsernameAndPassword(username, password);
}

测试发现在Dao层的接口进行注解时,拦截会在获取连接后执行,导致失效。具体原因没去深追,后续有时间再断点调试查看是什么原因,orm使用的是Sprng Data Jpa,使用mybatis貌似不会出现这种情况。
controler添加测试方法

1
2
3
4
5
6
7
8
@PostMapping("/test")
public AjaxResponseVO test(String username,String password) {
User userForLogin = service.findUserForLogin(username, password);
if (userForLogin != null) {
return AjaxResponseVO.success("登录成功");
}
return AjaxResponseVO.fail("用户名或密码错误");
}

post测试:

测试通过,打印日志: