cn.mybatis.mp.routing.datasource.SpringRoutingDataSource Maven / Gradle / Ivy
package cn.mybatis.mp.routing.datasource;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import javax.sql.DataSource;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.*;
/**
* 基于spring的动态数据源
*/
public class SpringRoutingDataSource extends AbstractRoutingDataSource {
private final Field resolvedDataSourcesField;
private final Field targetDataSourcesField;
public SpringRoutingDataSource() {
try {
resolvedDataSourcesField = AbstractRoutingDataSource.class.getDeclaredField("resolvedDataSources");
resolvedDataSourcesField.setAccessible(true);
targetDataSourcesField = AbstractRoutingDataSource.class.getDeclaredField("targetDataSources");
targetDataSourcesField.setAccessible(true);
} catch (NoSuchFieldException e) {
throw new RuntimeException(e);
}
}
@Override
protected String determineCurrentLookupKey() {
return DataSourceHolder.getCurrent();
}
@Override
protected DataSource determineTargetDataSource() {
DataSource dataSource = super.determineTargetDataSource();
if (dataSource instanceof GroupDataSource) {
return ((GroupDataSource) dataSource).loadBalance();
}
return dataSource;
}
public String getUrl() throws NoSuchMethodException {
DataSource dataSource = this.determineTargetDataSource();
String[] methodNames = new String[]{"getUrl", "getJdbcUrl"};
Throwable throwable = null;
for (String methodName : methodNames) {
try {
Method method = dataSource.getClass().getMethod(methodName);
return (String) method.invoke(dataSource);
} catch (Exception e) {
if (Objects.isNull(throwable)) {
throwable = e;
}
}
}
throw new NoSuchMethodException();
}
private String getGroupName(String key) {
if (!key.contains(Config.GROUP_SPLIT)) {
return null;
}
int index = key.lastIndexOf(Config.GROUP_SPLIT);
String no = key.substring(index + 1);
try {
Integer.parseInt(no);
} catch (NumberFormatException e) {
return null;
}
return key.substring(0, index);
}
/**
* 添加新的数据源
*
* @param key
* @param dataSource
*/
public void addNewDatasource(String key, DataSource dataSource) {
Map dataSourceMap = this.getSelfResolvedDataSources();
if (dataSourceMap.containsKey(key)) {
throw new DataSourceExistsException();
}
String groupName = getGroupName(key);
if (Objects.nonNull(groupName)) {
GroupDataSource groupDataSource = (GroupDataSource) dataSourceMap.get(key);
if (Objects.isNull(groupDataSource)) {
List list = new ArrayList<>();
list.add(dataSource);
groupDataSource = new GroupDataSource(list);
dataSourceMap.put(groupName, groupDataSource);
} else {
groupDataSource.getDelegateList().add(groupDataSource);
}
}
dataSourceMap.put(key, dataSource);
}
private Map getSelfResolvedDataSources() {
try {
return (Map) resolvedDataSourcesField.get(this);
} catch (IllegalAccessException e) {
throw new RuntimeException(e);
}
}
private Map getSelfTargetDataSources() {
try {
return (Map) this.targetDataSourcesField.get(this);
} catch (IllegalAccessException e) {
throw new RuntimeException(e);
}
}
/**
* 移除数据源
* 可移除分组数据源
*
* @param key
* @return 返回被移除的数据库
*/
public DataSource removeDatasource(String key) {
Map resolvedDataSources = this.getSelfResolvedDataSources();
Map targetDataSources = this.getSelfTargetDataSources();
DataSource dataSource = removeTargetByKey(key, resolvedDataSources, targetDataSources);
if (Objects.isNull(dataSource)) {
return null;
}
if (dataSource instanceof GroupDataSource) {
//移除分组数据源
GroupDataSource groupDataSource = (GroupDataSource) dataSource;
Iterator> iterable = resolvedDataSources.entrySet().iterator();
while (iterable.hasNext()) {
Map.Entry entry = iterable.next();
if (groupDataSource.getDelegateList().contains(entry.getValue())) {
removeTargetByKey(entry.getKey().toString(), targetDataSources);
iterable.remove();
}
}
} else {
String groupName = getGroupName(key);
if (Objects.nonNull(groupName)) {
GroupDataSource groupDataSource = (GroupDataSource) resolvedDataSources.get(groupName);
if (Objects.nonNull(groupDataSource)) {
List list = groupDataSource.getDelegateList();
list.remove(dataSource);
if (list.isEmpty()) {
//全部移除了 移除分组
removeTargetByKey(groupName, resolvedDataSources, targetDataSources);
}
}
}
}
return dataSource;
}
@SafeVarargs
private final DataSource removeTargetByKey(String key, Map... dataSourceMaps) {
DataSource dataSource = null;
for (Map dataSourceMap : dataSourceMaps) {
if (Objects.isNull(dataSource)) {
dataSource = dataSourceMap.remove(key);
} else {
dataSourceMap.remove(key);
}
}
return dataSource;
}
}