Spring Boot 2.X 设置动态数据源

1 多数据源实现的原理（AbstractRoutingDataSource）

Spring Boot 提供了抽象类 AbstractRoutingDataSource，通过扩展这个类实现根据不同的请求切换数据源。
AbstractRoutingDataSource继承AbstractDataSource，如果声明一个类DynamicDataSource继承AbstractRoutingDataSource后，DynamicDataSource本身就相当于一种数据源。

1.1 AbstractRoutingDataSource

public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean

AbstractRoutingDataSource 实现了 InitializingBean 那么spring在初始化该bean时，会调用AbstractRoutingDataSource 重写的afterPropertiesSet方法。
AbstractRoutingDataSource 继承了 AbstractDataSource，重写了getConnectionIf方法。

1.2 AbstractRoutingDataSource包含的属性

// 外部设置的多数据源容器
@Nullable  
private Map<Object, Object> targetDataSources;

// 外部设置默认的数据源
@Nullable
private Object defaultTargetDataSource;

private boolean lenientFallback = true;

private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();

// 内部的多数据源容器
@Nullable
private Map<Object, DataSource> resolvedDataSources;

// 内部的默认数据源
@Nullable
private DataSource resolvedDefaultDataSource;

1.3 主要方法解析

// 设置多数据源
public void setTargetDataSources(Map<Object, Object> targetDataSources) {
	this.targetDataSources = targetDataSources;
}

// 设置默认数据源
public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
	this.defaultTargetDataSource = defaultTargetDataSource;
}

// 重写getConnection方法，调用determineTargetDataSource返回DataSource的getConnection
@Override
public Connection getConnection() throws SQLException {
	return determineTargetDataSource().getConnection();
}

// 重写getConnection方法，调用determineTargetDataSource返回DataSource的getConnection
@Override
public Connection getConnection(String username, String password) throws SQLException {
	return determineTargetDataSource().getConnection(username, password);
}


// 初始化bean的时候，调用afterPropertiesSet
@Override
public void afterPropertiesSet() {
 	// 多数据源的map不可以为空，否则抛异常
	if (this.targetDataSources == null) {
		throw new IllegalArgumentException("Property 'targetDataSources' is required");
	}
	
	// 将外界传入的targetDataSources转到本类的resolvedDataSources 
	this.resolvedDataSources = new HashMap<>(this.targetDataSources.size());
	this.targetDataSources.forEach((key, value) -> {
		// 对key进行操作，转换为符合自己要求的key
		Object lookupKey = resolveSpecifiedLookupKey(key);
		// 获取每个key对应的DataSourced对象
		DataSource dataSource = resolveSpecifiedDataSource(value);
		// 添加到resolvedDataSources
		this.resolvedDataSources.put(lookupKey, dataSource);
	});
	if (this.defaultTargetDataSource != null) {
		// 将默认的数据源复制到本类里面的resolvedDefaultDataSource 
		this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
	}
}

// 子类可以重写该方法，默认方法只是简单的返回传入的值
protected Object resolveSpecifiedLookupKey(Object lookupKey) {
	return lookupKey;
}

// 根绝指定的数据源，转化为dataSource实例
protected DataSource resolveSpecifiedDataSource(Object dataSource) throws IllegalArgumentException {
	if (dataSource instanceof DataSource) {
		return (DataSource) dataSource;
	}
	else if (dataSource instanceof String) {
		//  注意：dataSourceLookup是JndiDataSourceLookup的实例
		return this.dataSourceLookup.getDataSource((String) dataSource);
	}
	else {
		throw new IllegalArgumentException(
				"Illegal data source value - only [javax.sql.DataSource] and String supported: " + dataSource);
	}
}

@Override
public Connection getConnection() throws SQLException {
	return determineTargetDataSource().getConnection();
}

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


// 见名知义，这个方法就是设定具体的数据源
protected DataSource determineTargetDataSource() {
	Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
	// 该方法需要自己实现
	Object lookupKey = determineCurrentLookupKey();
	// 获取要使用的数据源
	DataSource dataSource = this.resolvedDataSources.get(lookupKey);
	if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
		// 如果dataSource为null，返回设置的默认数据源
		dataSource = this.resolvedDefaultDataSource;
	}
	if (dataSource == null) {
		throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
	}
	return dataSource;
}

// 需要子类自己实现（必须），决定数据源的key是哪一个
@Nullable
protected abstract Object determineCurrentLookupKey();

2 多数据源实现（AOP + 注解）

/**
 * 动态数据源注解
 * @author heiky
 */
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface DataSource {

    /**
     *  数据源key值
     */
    String value() default "";

}

/**
 * 动态数据源上下文持有者
 * @author heiky
 */
public class DynamicDataSourceContextHolder {

    private static ThreadLocal<String> contextHolder= new ThreadLocal<>();

    /**
     * 切换数据源
     * @param key
     */
    public static void setDataSourceKey(String key) {
        contextHolder.set(key);
    }

    /**
     * 获取数据源
     * @return
     */
    public static String getDataSourceKey() {
        return contextHolder.get();
    }

    /**
     * 重置数据源
     */
    public static void clearDataSourceKey() {
        contextHolder.remove();
    }
}

/**
 * 自定义动态数据源
 * @author heiky
 */
public class DynamicDataSource extends AbstractRoutingDataSource {
 
    /**
     * 如果希望所有数据源在启动配置时就加载好，通过设置数据源Key值来切换数据
     */
    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.getDataSourceKey();
    }
}

/**
 * 动态数据源切换处理器（设置切面）
 * 切入点表达式用了两种，可以自行去了解
 *
 */
@Aspect
@Order(-1)  // 该切面应当先于 @Transactional 执行
@Component
@Slf4j
public class DynamicDataSourceAspect {

    @Pointcut("@annotation(cn.bigdata.customer.ds.DataSource)")
    public void dynamicDataSourceAspect() {
    }

    /**
     * 切换数据源
     *
     * @param point
     */
    @Before("dynamicDataSourceAspect()")
    public void switchDataSource(JoinPoint point) {
        MethodSignature signature = (MethodSignature) point.getSignature();
        DataSource dataSource = signature.getMethod().getAnnotation(DataSource.class);
        if (!DynamicDataSourceContextHolder.containDataSourceKey(dataSource.value())) {
            log.info("DataSource [{}] doesn't exist, use default DataSource [{}] " + dataSource.value());
        } else {
            // 切换数据源
            DynamicDataSourceContextHolder.setDataSourceKey(dataSource.value());
            log.info("Switch DataSource to [" + DynamicDataSourceContextHolder.getDataSourceKey()
                    + "] in Method [" + point.getSignature() + "]");
        }
    }

    /**
     * 重置数据源
     *
     * @param point
     * @param dataSource
     */
    @After("@annotation(dataSource)")
    public void restoreDataSource(JoinPoint point, DataSource dataSource) {
        // 将数据源置为默认数据源
        DynamicDataSourceContextHolder.clearDataSourceKey();
        log.info("Restore DataSource to [" + DynamicDataSourceContextHolder.getDataSourceKey()
                + "] in Method [" + point.getSignature() + "]");
    }
}


/**
 * 动态数据源的配置
 */
@Configuration
public class DataSourceConfig {

    @Bean("outernet")
    @ConfigurationProperties(prefix = "outernet")
    public DataSource outernet() {
        return DruidDataSourceBuilder.create().build();
    }

    @Bean("intranet")
    @ConfigurationProperties(prefix = "intranet")
    public DataSource intranet() {
        return DruidDataSourceBuilder.create().build();
    }

    /**
     * 注意:要在动态数据源+@primary
     * @return
     */
    @Bean("dynamicDataSource")
    @Primary
    public DataSource dynamicDataSource(DataSource outernet,DataSource intranet) {
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        Map<Object, Object> dataSourceMap = new HashMap<>(2);
        dataSourceMap.put("outernet", outernet);
        dataSourceMap.put("intranet", intranet);
        // 将 outernet 数据源作为默认指定的数据源
        dynamicDataSource.setDefaultDataSource(outernet);
        // 将 outernet 和 intranet 数据源作为指定的数据源
        dynamicDataSource.setDataSources(dataSourceMap);
        return dynamicDataSource;
    }

	 /**
     * 设置事务管理器
     * @return 
     */
	@Bean
	public PlatformTransactionManager transactionManager() {
	     // 配置事务管理, 使用事务时在方法头部添加@Transactional注解即可
	     return new DataSourceTransactionManager(dynamicDataSource());
	 }
}

下面是关于数据源的配置，可以自由配置，下面是笔者的配置方式

mybatis-plus.mapper-locations=classpath:/mapper/*Mapper.xml

outernet.type=com.alibaba.pool.DruidDataSource
outernet.driver-class-name=com.mysql.jdbc.Driver
outernet.url=jdbc:mysql://192.168.0.11:3306/consumer?useUnicode=true&characterEncoding=utf8&useSSL=false
outernet.username=root
outernet.password=123456
outernet.initial-size=10
outernet.min-idle=10
outernet.max-active=50
outernet.max-wait=2000
outernet.validation-query=SELECT 1
outernet.validation-query-timeout=2000
outernet.test-while-idle=true
outernet.test-on-borrow=false
outernet.test-on-return=false

intranet.type=com.alibaba.pool.DruidDataSource
intranet.driver-class-name=com.mysql.jdbc.Driver
intranet.url=jdbc:mysql://192.168.0.12:3306/producer?useUnicode=true&characterEncoding=utf8&useSSL=false
intranet.username=root
intranet.password=123456
intranet.initial-size=10
intranet.min-idle=10
intranet.max-active=50
intranet.max-wait=2000
intranet.validation-query=SELECT 1
intranet.validation-query-timeout=2000
intranet.test-while-idle=true
intranet.test-on-borrow=false
intranet.test-on-return=false

3 注意事项

使用AbstractRoutingDataSource来管理数据源的话，会有一个问题：事务和切换数据源不能同时生效（在同一包下面）。
为什么会有这个问题呢？
因为一般我们的事务是以AOP的方式，添加注解在我们业务方法上，并且是通过获取Connection对象来实现的（这应该是事务的通用方式），事务本身也包含数据源，恰好和从AbstractRoutingDataSource对象中切换数据源的时机重合，就是调determineTargetDataSource方法，这也是我们实现AbstractRoutingDataSource类唯一要覆盖的方法。所以在连接池的环境下，就导致由事务创建了链接，也是调用了determineTargetDataSource方法。如果我们要向切换数据源就必须要在事务之前完成，但这就无法实现在事务中使用多数据源了。
解决方式：分包，在调用有事务的方法事前切换数据源

4 数据源切换的整体流程

细心的读者，可能会有个疑问，什么时候会调用自定义的DynamicDataSource的getConnection方法？
因为dynamicDataSource在Spring容器里面定义成了一个bean，容器里面就这一个数据源（必须用到的）。Mybatis是连接Spring与数据库的中间件，其中SqlSessionFactory的openSession就是getConnection的实现。