實現(xiàn)效果

先說效果，要實現(xiàn)方法級別注解切換當(dāng)前數(shù)據(jù)源，不設(shè)置注解時走默認(rèn)數(shù)據(jù)源，同時支持JNDI源。

總體思路

Spring框架中存在一個抽象類AbstractRoutingDataSource，他是一個可以動態(tài)選擇當(dāng)前DataSource的路由類，我們就是要從這里入手，重新實現(xiàn)數(shù)據(jù)源的切換選擇邏輯。然后借助注解和切面，將當(dāng)前需要的數(shù)據(jù)源名稱放在ThreadLocal中，需要時從當(dāng)前線程取得即可完成數(shù)據(jù)源的切換。
注解部分比較簡單不再詳說，看AbstractRoutingDataSource。該類文檔寫的非常全面，自行翻譯一下就可以看懂。主要看其中的幾個關(guān)鍵方法。

setTargetDataSources

類中存在一個成員變量targetDataSources，結(jié)合之后的setTargetDataSources方法可知，這里用來保存目標(biāo)數(shù)據(jù)源。
根據(jù)注釋我們可以知道，targetDataSources的key可以是數(shù)據(jù)源的名字，value是相應(yīng)數(shù)據(jù)源的實例。
當(dāng)然這里也可是使用其他的保存方式，然后自行改寫用來查找數(shù)據(jù)源的determineCurrentLookupKey方法，默認(rèn)場景就足夠我們使用了。所以我們要構(gòu)建一個Map出來，其中key用來區(qū)分?jǐn)?shù)據(jù)源的名字，value放入對應(yīng)數(shù)據(jù)源的實例，有幾個數(shù)據(jù)源就放幾個進去。

	@Nullableprivate Map<Object, Object> targetDataSources;/*** Specify the map of target DataSources, with the lookup key as key.* The mapped value can either be a corresponding {@link javax.sql.DataSource}* instance or a data source name String (to be resolved via a* {@link #setDataSourceLookup DataSourceLookup}).* <p>The key can be of arbitrary type; this class implements the* generic lookup process only. The concrete key representation will* be handled by {@link #resolveSpecifiedLookupKey(Object)} and* {@link #determineCurrentLookupKey()}.*/public void setTargetDataSources(Map<Object, Object> targetDataSources) {this.targetDataSources = targetDataSources;}

setDefaultTargetDataSource

上面說了如何設(shè)置當(dāng)前數(shù)據(jù)源，那如果在開發(fā)的時候每一個方法都要聲明一下使用哪個源就太麻煩了，所以Spring提供了一個方法用來設(shè)置默認(rèn)的數(shù)據(jù)源，沒啥可說的，傳入DataSource實例就好了。

	@Nullableprivate Object defaultTargetDataSource;/*** Specify the default target DataSource, if any.* <p>The mapped value can either be a corresponding {@link javax.sql.DataSource}* instance or a data source name String (to be resolved via a* {@link #setDataSourceLookup DataSourceLookup}).* <p>This DataSource will be used as target if none of the keyed* {@link #setTargetDataSources targetDataSources} match the* {@link #determineCurrentLookupKey()} current lookup key.*/public void setDefaultTargetDataSource(Object defaultTargetDataSource) {this.defaultTargetDataSource = defaultTargetDataSource;}

determineCurrentLookupKey

在設(shè)置好數(shù)據(jù)源之后，接下來這幾個尋路方法則是能實現(xiàn)動態(tài)數(shù)據(jù)源切換的重點。afterPropertiesSet方法對我們以配置的數(shù)據(jù)源進行校驗；如果我們在第一步配置數(shù)據(jù)源map的時候?qū)ey有特殊處理則要自己實現(xiàn)抽象方法resolveSpecifiedLookupKey，告訴Spring應(yīng)該怎么解析這個key值；determineTargetDataSource則最終確定要使用哪一個數(shù)據(jù)源，其中有一個方法determineCurrentLookupKey需要關(guān)注，這個方法會返回當(dāng)前要使用的數(shù)據(jù)源名字，但他是個抽象方法，所以我們需要給他重寫一下，改為從當(dāng)前線程獲取數(shù)據(jù)源名稱。

	@Overridepublic 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);}}@Overridepublic Connection getConnection(String username, String password) throws SQLException {return determineTargetDataSource().getConnection(username, password);}/*** Retrieve the current target DataSource. Determines the* {@link #determineCurrentLookupKey() current lookup key}, performs* a lookup in the {@link #setTargetDataSources targetDataSources} map,* falls back to the specified* {@link #setDefaultTargetDataSource default target DataSource} if necessary.* @see #determineCurrentLookupKey()*/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 = this.resolvedDefaultDataSource;}if (dataSource == null) {throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");}return dataSource;}/*** Determine the current lookup key. This will typically be* implemented to check a thread-bound transaction context.* <p>Allows for arbitrary keys. The returned key needs* to match the stored lookup key type, as resolved by the* {@link #resolveSpecifiedLookupKey} method.*/@Nullableprotected abstract Object determineCurrentLookupKey();

代碼實現(xiàn)

思路理順了，代碼寫起來就比較快，直接貼最后代碼，部分地方保留了注釋。

數(shù)據(jù)源切換注解

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;/*** 數(shù)據(jù)源切換注解，默認(rèn)為primary*/
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface TargetDataSource {DataSourceEnum value() default DataSourceEnum.PRIMARY;
}

數(shù)據(jù)源切換切面

這里需要特別提醒一下，事務(wù)注解@Transactional默認(rèn)處于切面代理的最后一個，所以我們需要保證數(shù)據(jù)源切換注解優(yōu)先級要高于事務(wù)注解。

import org.aspectj.lang.annotation.After;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;/*** 數(shù)據(jù)源切換切面*/
@Aspect
@Component
@Order(1)
public class DynamicDataSourceAspect {private final static Logger log = LoggerFactory.getLogger(DynamicDataSourceAspect.class);@Before(value = "@annotation(targetDataSource)")public void beforePointCut(TargetDataSource targetDataSource) {log.debug("數(shù)據(jù)源切換為 " + targetDataSource.value().getDataSourceName());DynamicDataSourceContextHolder.setDataSource(targetDataSource.value().getDataSourceName());}@After(value = "@annotation(targetDataSource)")public void afterPointCut(TargetDataSource targetDataSource) {log.debug("數(shù)據(jù)源恢復(fù)為 " + DataSourceEnum.PRIMARY.getDataSourceName());DynamicDataSourceContextHolder.clearDataSource();}
}

數(shù)據(jù)源枚舉類

/*** 數(shù)據(jù)源枚舉類*/
public enum DataSourceEnum {PRIMARY("primary"), SECONDARY("secondary");private final String dataSourceName;public String getDataSourceName() {return dataSourceName;}DataSourceEnum(String dataSourceName) {this.dataSourceName = dataSourceName;}
}

數(shù)據(jù)源上下文保持類

/*** 數(shù)據(jù)源上下文線程持有類*/
public class DynamicDataSourceContextHolder {/*** 存放當(dāng)前線程使用的數(shù)據(jù)源類型信息*/private static final ThreadLocal<String> CONTEXT_HOLDER = new ThreadLocal<>();public static void setDataSource(String dataSourceType) {CONTEXT_HOLDER.set(dataSourceType);}public static String getDataSource() {return CONTEXT_HOLDER.get();}public static void clearDataSource() {CONTEXT_HOLDER.remove();}
}

AbstractRoutingDataSource自定義實現(xiàn)

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;import javax.sql.DataSource;
import java.util.Map;/*** 動態(tài)數(shù)據(jù)源切換類** @author liuenqi*/
public class DynamicDataSource extends AbstractRoutingDataSource {@Overrideprotected Object determineCurrentLookupKey() {return DynamicDataSourceContextHolder.getDataSource();}public DynamicDataSource(DataSource defaultTargetDataSource, Map<Object, Object> targetDataSources) {// 默認(rèn)數(shù)據(jù)源super.setDefaultTargetDataSource(defaultTargetDataSource);// 所有目標(biāo)數(shù)據(jù)源super.setTargetDataSources(targetDataSources);// 后處理super.afterPropertiesSet();}
}

數(shù)據(jù)源注冊

注意使用jndi源的時候需要加一個特定前綴。

import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.config.ConstructorArgumentValues;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.env.Environment;
import org.springframework.core.type.AnnotationMetadata;import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;/*** 多數(shù)據(jù)源注冊類*/
public class DynamicDataSourceRegister implements ImportBeanDefinitionRegistrar, EnvironmentAware {private DataSource primaryDataSource;private DataSource secondaryDataSource;@Overridepublic void setEnvironment(Environment environment) {initPrimaryDataSource(environment);initSecondaryDataSource(environment);}/*** 組裝主數(shù)據(jù)源參數(shù)，兼容jdbc-url與jndi** @param env Environment*/private void initPrimaryDataSource(Environment env) {Map<String, String> paramMap = new HashMap<>(4);if (StringUtils.isNotBlank(env.getProperty("spring.datasource.primary.url"))) {paramMap.put("url", env.getProperty("spring.datasource.primary.url"));paramMap.put("userName", env.getProperty("spring.datasource.primary.username"));paramMap.put("password", env.getProperty("spring.datasource.primary.password"));paramMap.put("driverClassName", env.getProperty("spring.datasource.primary.driver-class-name"));} else {paramMap.put("jndi", env.getProperty("spring.datasource.primary.jndi-name"));}primaryDataSource = buildDataSource(paramMap);}/*** 組裝輔數(shù)據(jù)源參數(shù)，兼容jdbc-url與jndi** @param env Environment*/private void initSecondaryDataSource(Environment env) {if (StringUtils.isNotBlank(env.getProperty("spring.datasource.secondary.url"))) {Map<String, String> paramMap = new HashMap<>(4);paramMap.put("url", env.getProperty("spring.datasource.secondary.url"));paramMap.put("userName", env.getProperty("spring.datasource.secondary.username"));paramMap.put("password", env.getProperty("spring.datasource.secondary.password"));paramMap.put("driverClassName", env.getProperty("spring.datasource.secondary.driver-class-name"));secondaryDataSource = buildDataSource(paramMap);} else if (StringUtils.isNotBlank(env.getProperty("spring.datasource.secondary.jndi-name"))) {Map<String, String> paramMap = new HashMap<>(2);paramMap.put("jndi", env.getProperty("spring.datasource.secondary.jndi-name"));secondaryDataSource = buildDataSource(paramMap);}}@Overridepublic void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry) {Map<Object, Object> targetDataSource = new HashMap<>(2);targetDataSource.put("primary", primaryDataSource);if (Objects.nonNull(secondaryDataSource)) {targetDataSource.put("secondary", secondaryDataSource);}// 為DynamicDataSource構(gòu)造參數(shù)，注意參數(shù)順序ConstructorArgumentValues constructorArgumentValues = new ConstructorArgumentValues();constructorArgumentValues.addGenericArgumentValue(primaryDataSource);constructorArgumentValues.addGenericArgumentValue(targetDataSource);// 構(gòu)造bean放入IOCGenericBeanDefinition beanDefinition = new GenericBeanDefinition();beanDefinition.setBeanClass(DynamicDataSource.class);beanDefinition.setConstructorArgumentValues(constructorArgumentValues);beanDefinition.setSynthetic(true);registry.registerBeanDefinition("dataSource", beanDefinition);}/*** 使用HikariDataSource** @param paramMap {"url":"JDBC-URL","userName":"數(shù)據(jù)庫用戶名","password":"密碼","driverClassName":"驅(qū)動名","jndi":"jndi源"}* @return HikariDataSource*/private DataSource buildDataSource(Map<String, String> paramMap) {HikariConfig hikariConfig = new HikariConfig();if (paramMap.containsKey("url")) {hikariConfig.setJdbcUrl(paramMap.get("url"));hikariConfig.setUsername(paramMap.get("userName"));hikariConfig.setPassword(paramMap.get("password"));hikariConfig.setDriverClassName(paramMap.get("driverClassName"));} else {hikariConfig.setDataSourceJNDI("java:comp/env/" + paramMap.get("jndi"));}return new HikariDataSource(hikariConfig);}
}

application.yml配置

spring:datasource:primary:url: jdbc:mysql://xxxxxusername: xxxxpassword: xxxxxdriver-class-name: com.mysql.cj.jdbc.Driversecondary: jndi-name: jdbc/db