基于注解切换、Hikari实现的SpringBoot动态数据源（支持JNDI）

实现效果

先说效果，要实现方法级别注解切换当前数据源，不设置注解时走默认数据源，同时支持JNDI源。

总体思路

Spring框架中存在一个抽象类AbstractRoutingDataSource，他是一个可以动态选择当前DataSource的路由类，我们就是要从这里入手，重新实现数据源的切换选择逻辑。然后借助注解和切面，将当前需要的数据源名称放在ThreadLocal中，需要时从当前线程取得即可完成数据源的切换。
注解部分比较简单不再详说，看AbstractRoutingDataSource。该类文档写的非常全面，自行翻译一下就可以看懂。主要看其中的几个关键方法。

setTargetDataSources

类中存在一个成员变量targetDataSources，结合之后的setTargetDataSources方法可知，这里用来保存目标数据源。
根据注释我们可以知道，targetDataSources的key可以是数据源的名字，value是相应数据源的实例。
当然这里也可是使用其他的保存方式，然后自行改写用来查找数据源的determineCurrentLookupKey方法，默认场景就足够我们使用了。所以我们要构建一个Map出来，其中key用来区分数据源的名字，value放入对应数据源的实例，有几个数据源就放几个进去。

	@Nullable
	private 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}).
	 * 
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

上面说了如何设置当前数据源，那如果在开发的时候每一个方法都要声明一下使用哪个源就太麻烦了，所以Spring提供了一个方法用来设置默认的数据源，没啥可说的，传入DataSource实例就好了。

	@Nullable
	private Object defaultTargetDataSource;

	/**
	 * Specify the default target DataSource, if any.
	 * 
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}).
	 * 
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

在设置好数据源之后，接下来这几个寻路方法则是能实现动态数据源切换的重点。afterPropertiesSet方法对我们以配置的数据源进行校验；如果我们在第一步配置数据源map的时候对key有特殊处理则要自己实现抽象方法resolveSpecifiedLookupKey，告诉Spring应该怎么解析这个key值；determineTargetDataSource则最终确定要使用哪一个数据源，其中有一个方法determineCurrentLookupKey需要关注，这个方法会返回当前要使用的数据源名字，但他是个抽象方法，所以我们需要给他重写一下，改为从当前线程获取数据源名称。

	@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);
		}
	}

	@Override
	public 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.
	 * 
Allows for arbitrary keys. The returned key needs
	 * to match the stored lookup key type, as resolved by the
	 * {@link #resolveSpecifiedLookupKey} method.
	 */
	@Nullable
	protected abstract Object determineCurrentLookupKey();

代码实现

思路理顺了，代码写起来就比较快，直接贴最后代码，部分地方保留了注释。

数据源切换注解

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * 数据源切换注解，默认为primary
 */
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface TargetDataSource {
    DataSourceEnum value() default DataSourceEnum.PRIMARY;
}

数据源切换切面

这里需要特别提醒一下，事务注解@Transactional默认处于切面代理的最后一个，所以我们需要保证数据源切换注解优先级要高于事务注解。

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;

/**
 * 数据源切换切面
 */
@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("数据源切换为 " + targetDataSource.value().getDataSourceName());
        DynamicDataSourceContextHolder.setDataSource(targetDataSource.value().getDataSourceName());
    }

    @After(value = "@annotation(targetDataSource)")
    public void afterPointCut(TargetDataSource targetDataSource) {
        log.debug("数据源恢复为 " + DataSourceEnum.PRIMARY.getDataSourceName());
        DynamicDataSourceContextHolder.clearDataSource();
    }
}

数据源枚举类

/**
 * 数据源枚举类
 */
public enum DataSourceEnum {

    PRIMARY("primary"), SECONDARY("secondary");

    private final String dataSourceName;

    public String getDataSourceName() {
        return dataSourceName;
    }

    DataSourceEnum(String dataSourceName) {
        this.dataSourceName = dataSourceName;
    }
}

数据源上下文保持类

/**
 * 数据源上下文线程持有类
 */
public class DynamicDataSourceContextHolder {

    /**
     * 存放当前线程使用的数据源类型信息
     */
    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自定义实现

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.Map;

/**
 * 动态数据源切换类
 *
 * @author liuenqi
 */
public class DynamicDataSource extends AbstractRoutingDataSource {
    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.getDataSource();
    }

    public DynamicDataSource(DataSource defaultTargetDataSource, Map<Object, Object> targetDataSources) {
        // 默认数据源
        super.setDefaultTargetDataSource(defaultTargetDataSource);
        // 所有目标数据源
        super.setTargetDataSources(targetDataSources);
        // 后处理
        super.afterPropertiesSet();
    }
}

数据源注册

注意使用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;

/**
 * 多数据源注册类
 */
public class DynamicDataSourceRegister implements ImportBeanDefinitionRegistrar, EnvironmentAware {

    private DataSource primaryDataSource;
    private DataSource secondaryDataSource;

    @Override
    public void setEnvironment(Environment environment) {
        initPrimaryDataSource(environment);
        initSecondaryDataSource(environment);
    }

    /**
     * 组装主数据源参数，兼容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);
    }

    /**
     * 组装辅数据源参数，兼容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);
        }
    }

    @Override
    public 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构造参数，注意参数顺序
        ConstructorArgumentValues constructorArgumentValues = new ConstructorArgumentValues();
        constructorArgumentValues.addGenericArgumentValue(primaryDataSource);
        constructorArgumentValues.addGenericArgumentValue(targetDataSource);

        // 构造bean放入IOC
        GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
        beanDefinition.setBeanClass(DynamicDataSource.class);
        beanDefinition.setConstructorArgumentValues(constructorArgumentValues);
        beanDefinition.setSynthetic(true);

        registry.registerBeanDefinition("dataSource", beanDefinition);
    }

    /**
     * 使用HikariDataSource
     *
     * @param paramMap {"url":"JDBC-URL","userName":"数据库用户名","password":"密码","driverClassName":"驱动名","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);
    }
}

启动类配置

@Import({DynamicDataSourceRegister.class})

application.yml配置

spring:
  datasource:
    primary:
      url: jdbc:mysql://xxxxx
      username: xxxx
      password: xxxxx
      driver-class-name: com.mysql.cj.jdbc.Driver
    secondary: 
      jndi-name: jdbc/db