dynamic-datasource源码分析

0. 动态切换数据源的实现原理

通过实现javax.sql.DataSource对连接进行动态切换管理;将数据源名称放在ThreadLocal中;通过数据源名称来构造对应的数据库连接进行实现

1. 整体架构

1.1 代码结构

1.2 整体结构

上面是支持的功能,下面是各个支持的组件

2. 源码分析(version:3.5.1)

首先是核心配置类

@Slf4j
@Configuration
@EnableConfigurationProperties(DynamicDataSourceProperties.class)
@AutoConfigureBefore(value = DataSourceAutoConfiguration.class, name = "com.alibaba.druid.spring.boot.autoconfigure.DruidDataSourceAutoConfigure")
@Import(value = {DruidDynamicDataSourceConfiguration.class, DynamicDataSourceCreatorAutoConfiguration.class})
@ConditionalOnProperty(prefix = DynamicDataSourceProperties.PREFIX, name = "enabled", havingValue = "true", matchIfMissing = true)
public class DynamicDataSourceAutoConfiguration implements InitializingBean {

	private final DynamicDataSourceProperties properties;

	/**
	* 可以自己实现DynamicDataSourcePropertiesCustomizer对配置文件进行定制
	*/
	private final List<DynamicDataSourcePropertiesCustomizer> dataSourcePropertiesCustomizers;

	public DynamicDataSourceAutoConfiguration(
		DynamicDataSourceProperties properties,
		ObjectProvider<List<DynamicDataSourcePropertiesCustomizer>> dataSourcePropertiesCustomizers) {
		this.properties = properties;
		this.dataSourcePropertiesCustomizers = dataSourcePropertiesCustomizers.getIfAvailable();
	}

	/**
	* 多数据源加载实现 启动默认从yml配置文件中进行读取
	* 可以自己实现加载方式
	* @return
	*/
	@Bean
	public DynamicDataSourceProvider ymlDynamicDataSourceProvider() {
		return new YmlDynamicDataSourceProvider(properties.getDatasource());
	}

	/**
	* 核心动态数据源组件DynamicRoutingDataSource
	* 提供了动态数据源切换的功能
	* @return DynamicRoutingDataSource
	*/
	@Bean
	@ConditionalOnMissingBean
	public DataSource dataSource() {
		DynamicRoutingDataSource dataSource = new DynamicRoutingDataSource();
		dataSource.setPrimary(properties.getPrimary());
		dataSource.setStrict(properties.getStrict());
		dataSource.setStrategy(properties.getStrategy());
		dataSource.setP6spy(properties.getP6spy());
		dataSource.setSeata(properties.getSeata());
		return dataSource;
	}

	/**
	* 动态多数据源注解通知dynamicDatasourceAnnotationAdvisor,
	* 这是一个aop前置通知,当一个请求发生的时候,会触发前置通知,通过注解的内容切换对应的数据源
	* @param dsProcessor
	* @return
	*/
	@Role(BeanDefinition.ROLE_INFRASTRUCTURE)
	@Bean
	@ConditionalOnProperty(prefix = DynamicDataSourceProperties.PREFIX + ".aop", name = "enabled", havingValue = "true", matchIfMissing = true)
	public Advisor dynamicDatasourceAnnotationAdvisor(DsProcessor dsProcessor) {
		DynamicDatasourceAopProperties aopProperties = properties.getAop();
		DynamicDataSourceAnnotationInterceptor interceptor = new DynamicDataSourceAnnotationInterceptor(aopProperties.getAllowedPublicOnly(), dsProcessor);
		DynamicDataSourceAnnotationAdvisor advisor = new DynamicDataSourceAnnotationAdvisor(interceptor, DS.class);
		advisor.setOrder(aopProperties.getOrder());
		return advisor;
	}

	/**
	* 多数据源事物拦截器DynamicLocalTransactionInterceptor
	* @return
	*/
	@Role(BeanDefinition.ROLE_INFRASTRUCTURE)
	@Bean
	@ConditionalOnProperty(prefix = DynamicDataSourceProperties.PREFIX, name = "seata", havingValue = "false", matchIfMissing = true)
	public Advisor dynamicTransactionAdvisor() {
		DynamicLocalTransactionInterceptor interceptor = new DynamicLocalTransactionInterceptor();
		return new DynamicDataSourceAnnotationAdvisor(interceptor, DSTransactional.class);
	}

	/**
	* 多数据源连接池创建事件
	* 目前用于数据库密码解密
	* @return
	*/
	@Bean
	@ConditionalOnMissingBean
	public DataSourceInitEvent dataSourceInitEvent() {
		return new EncDataSourceInitEvent();
	}

	/**
	* 动态多数据源预置处理器dsProcess,ds就是datasource的简称。这里主要采用的是责任链设计模式,获取ds。
	* @param beanFactory
	* @return
	*/
	@Bean
	@ConditionalOnMissingBean
	public DsProcessor dsProcessor(BeanFactory beanFactory) {
		DsHeaderProcessor headerProcessor = new DsHeaderProcessor();
		DsSessionProcessor sessionProcessor = new DsSessionProcessor();
		DsSpelExpressionProcessor spelExpressionProcessor = new DsSpelExpressionProcessor();
		spelExpressionProcessor.setBeanResolver(new BeanFactoryResolver(beanFactory));
		headerProcessor.setNextProcessor(sessionProcessor);
		sessionProcessor.setNextProcessor(spelExpressionProcessor);
		return headerProcessor;
	}

	@Override
	public void afterPropertiesSet() {
		if (!CollectionUtils.isEmpty(dataSourcePropertiesCustomizers)) {
			for (DynamicDataSourcePropertiesCustomizer customizer : dataSourcePropertiesCustomizers) {
				customizer.customize(properties);
			}
		}
	}
}
复制代码

基于拦截器来获取数据源

  • 通过注解来解析对应的数据源
/**
 * Core Interceptor of Dynamic Datasource
 *
 * @author TaoYu
 * @since 1.2.0
 */
public class DynamicDataSourceAnnotationInterceptor implements MethodInterceptor {

    /**
     * The identification of SPEL.
     */
    private static final String DYNAMIC_PREFIX = "#";

    private final DataSourceClassResolver dataSourceClassResolver;
    private final DsProcessor dsProcessor;

    public DynamicDataSourceAnnotationInterceptor(Boolean allowedPublicOnly, DsProcessor dsProcessor) {
        dataSourceClassResolver = new DataSourceClassResolver(allowedPublicOnly);
        this.dsProcessor = dsProcessor;
    }

    @Override
    public Object invoke(MethodInvocation invocation) throws Throwable {
        String dsKey = determineDatasourceKey(invocation);
        DynamicDataSourceContextHolder.push(dsKey);
        try {
            return invocation.proceed();
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    /**
     * 判断注解是否以#开头 如果#开头 则通过dsProcessor责任链来查询对应的数据源
     * 否则直接以注解内容的数据源名称,若注解数据源不存在,则使用默认数据源
     * @param invocation
     * @return
     */
    private String determineDatasourceKey(MethodInvocation invocation) {
        String key = dataSourceClassResolver.findKey(invocation.getMethod(), invocation.getThis());
        return key.startsWith(DYNAMIC_PREFIX) ? dsProcessor.determineDatasource(invocation, key) : key;
    }
}
复制代码

使用DsProcess来获取数据源

    /**
     * 动态多数据源预置处理器dsProcess,ds就是datasource的简称。这里主要采用的是责任链设计模式,获取ds。
     * 首先在header中获取,header中没有,去session中获取, session中也没有, 通过spel获取。
     * @param beanFactory
     * @return
     */
    @Bean
    @ConditionalOnMissingBean
    public DsProcessor dsProcessor(BeanFactory beanFactory) {
        DsHeaderProcessor headerProcessor = new DsHeaderProcessor();
        DsSessionProcessor sessionProcessor = new DsSessionProcessor();
        DsSpelExpressionProcessor spelExpressionProcessor = new DsSpelExpressionProcessor();
        spelExpressionProcessor.setBeanResolver(new BeanFactoryResolver(beanFactory));
        headerProcessor.setNextProcessor(sessionProcessor);
        sessionProcessor.setNextProcessor(spelExpressionProcessor);
        return headerProcessor;
    }
复制代码
public abstract class DsProcessor {

    private DsProcessor nextProcessor;

    public void setNextProcessor(DsProcessor dsProcessor) {
        this.nextProcessor = dsProcessor;
    }

    /**
     * 抽象匹配条件 匹配才会走当前执行器否则走下一级执行器
     *
     * @param key DS注解里的内容
     * @return 是否匹配
     */
    public abstract boolean matches(String key);

    /**
     * 决定数据源
     * <pre>
     *     调用底层doDetermineDatasource,
     *     如果返回的是null则继续执行下一个,否则直接返回
     * </pre>
     *
     * @param invocation 方法执行信息
     * @param key        DS注解里的内容
     * @return 数据源名称
     */
    public String determineDatasource(MethodInvocation invocation, String key) {
        if (matches(key)) {
            String datasource = doDetermineDatasource(invocation, key);
            if (datasource == null && nextProcessor != null) {
                return nextProcessor.determineDatasource(invocation, key);
            }
            return datasource;
        }
        if (nextProcessor != null) {
            return nextProcessor.determineDatasource(invocation, key);
        }
        return null;
    }

    /**
     * 抽象最终决定数据源
     *
     * @param invocation 方法执行信息
     * @param key        DS注解里的内容
     * @return 数据源名称
     */
    public abstract String doDetermineDatasource(MethodInvocation invocation, String key);
}
复制代码
public class DsHeaderProcessor extends DsProcessor {

    /**
     * header prefix
     */
    private static final String HEADER_PREFIX = "#header";

    @Override
    public boolean matches(String key) {
        return key.startsWith(HEADER_PREFIX);
    }

    /**
     * 从request.getHeader中获取数据源信息
     * @param invocation 方法执行信息
     * @param key        DS注解里的内容
     * @return
     */
    @Override
    public String doDetermineDatasource(MethodInvocation invocation, String key) {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        return request.getHeader(key.substring(8));
    }
}
复制代码
public class DsSpelExpressionProcessor extends DsProcessor {

    /**
     * 参数发现器
     */
    private static final ParameterNameDiscoverer NAME_DISCOVERER = new DefaultParameterNameDiscoverer();
    /**
     * Express语法解析器
     */
    private static final ExpressionParser PARSER = new SpelExpressionParser();
    /**
     * 解析上下文的模板
     * 对于默认不设置的情况下,从参数中取值的方式 #param1
     * 设置指定模板 ParserContext.TEMPLATE_EXPRESSION 后的取值方式: #{#param1}
     * issues: https://github.com/baomidou/dynamic-datasource-spring-boot-starter/issues/199
     */
    private ParserContext parserContext = new ParserContext() {

        @Override
        public boolean isTemplate() {
            return false;
        }

        @Override
        public String getExpressionPrefix() {
            return null;
        }

        @Override
        public String getExpressionSuffix() {
            return null;
        }
    };
    private BeanResolver beanResolver;

    @Override
    public boolean matches(String key) {
        return true;
    }

    @Override
    public String doDetermineDatasource(MethodInvocation invocation, String key) {
        Method method = invocation.getMethod();
        Object[] arguments = invocation.getArguments();
        StandardEvaluationContext context = new MethodBasedEvaluationContext(null, method, arguments, NAME_DISCOVERER);
        context.setBeanResolver(beanResolver);
        final Object value = PARSER.parseExpression(key, parserContext).getValue(context);
        return value == null ? null : value.toString();
    }

    public void setParserContext(ParserContext parserContext) {
        this.parserContext = parserContext;
    }

    public void setBeanResolver(BeanResolver beanResolver) {
        this.beanResolver = beanResolver;
    }
}
复制代码
  • 通过指定某个方法使用某个数据源
@Slf4j
public class DynamicDatasourceNamedInterceptor implements MethodInterceptor {

	private static final String DYNAMIC_PREFIX = "#";
	private final Map<String, String> nameMap = new HashMap<>();
	private final DsProcessor dsProcessor;

	public DynamicDatasourceNamedInterceptor(DsProcessor dsProcessor) {
		this.dsProcessor = dsProcessor;
	}

	@Nullable
	@Override
	public Object invoke(@Nonnull MethodInvocation invocation) throws Throwable {
		String dsKey = determineDatasourceKey(invocation);
		DynamicDataSourceContextHolder.push(dsKey);
		try {
			return invocation.proceed();
		} finally {
			DynamicDataSourceContextHolder.poll();
		}
	}

	/**
     * add Item Pattern
     *
     * @param methodName like select*
     * @param dsKey      like master or slave
     */
	public void addPattern(@Nonnull String methodName, @Nonnull String dsKey) {
		log.debug("dynamic-datasource adding ds method [" + methodName + "] with attribute [" + dsKey + "]");
		nameMap.put(methodName, dsKey);
	}

	/**
     * add PatternMap
     *
     * @param map namedMap
     */
	public void addPatternMap(Map<String, String> map) {
		for (Map.Entry<String, String> item : map.entrySet()) {
			addPattern(item.getKey(), item.getValue());
		}
	}

	/**
     * config from properties
     * <pre>
     *         Properties attributes = new Properties();
     *         attributes.setProperty("select*", "slave");
     *         attributes.setProperty("add*", "master");
     *         attributes.setProperty("update*", "master");
     *         attributes.setProperty("delete*", "master");
     * </pre>
     *
     * @param properties ds properties
     */
	public void fromProperties(@Nonnull Properties properties) {
		Enumeration<?> propNames = properties.propertyNames();
		while (propNames.hasMoreElements()) {
			String methodName = (String) propNames.nextElement();
			String value = properties.getProperty(methodName);
			this.addPattern(methodName, value);
		}
	}


	private boolean isMatch(String methodName, String mappedName) {
		return PatternMatchUtils.simpleMatch(mappedName, methodName);
	}

	private String determineDatasourceKey(MethodInvocation invocation) {
		String key = findDsKey(invocation);
		return (key != null && key.startsWith(DYNAMIC_PREFIX)) ? dsProcessor.determineDatasource(invocation, key) : key;
	}

	private String findDsKey(MethodInvocation invocation) {
		Method method = invocation.getMethod();
		if (!ClassUtils.isUserLevelMethod(method)) {
			return null;
		}

		// Look for direct name match.
		String methodName = method.getName();
		String dsKey = this.nameMap.get(methodName);

		if (dsKey == null) {
			// Look for most specific name match.
			String bestNameMatch = null;
			for (String mappedName : this.nameMap.keySet()) {
				boolean match1 = isMatch(methodName, mappedName);
				boolean match2 = bestNameMatch == null || bestNameMatch.length() <= mappedName.length();
				if (match1 && match2) {
					dsKey = this.nameMap.get(mappedName);
					bestNameMatch = mappedName;
				}
			}
		}
		return dsKey;
	}
}
复制代码

数据源创建者

数据源创建等底层操作,上层根据使用的不同类型的数据源进行动态数据源的创建;代码结构如下

由一个接口DataSourceCreator,一个抽象类AbstractDataSourceCreator以及若干个不同类型的数据源实现类组成.

public interface DataSourceCreator {

    /**
     * 通过属性创建数据源
     *
     * @param dataSourceProperty 数据源属性
     * @return 被创建的数据源
     */
    DataSource createDataSource(DataSourceProperty dataSourceProperty);

    /**
     * 当前创建器是否支持根据此属性创建
     * 
     * @param dataSourceProperty 数据源属性
     * @return 是否支持
     */
    boolean support(DataSourceProperty dataSourceProperty);
}
复制代码
/**
 * 抽象连接池创建器
 * <p>
 * 这里主要处理一些公共逻辑,如脚本和事件等
 *
 * @author TaoYu
 */
@Slf4j
public abstract class AbstractDataSourceCreator implements DataSourceCreator {

    @Autowired
    protected DynamicDataSourceProperties properties;
    @Autowired
    protected DataSourceInitEvent dataSourceInitEvent;

    /**
     * 子类去实际创建连接池
     *
     * @param dataSourceProperty 数据源信息
     * @return 实际连接池
     */
    public abstract DataSource doCreateDataSource(DataSourceProperty dataSourceProperty);

    /**
     * 创建数据源
     * 1.创建之前数据库密码解密
     * 2.创建数据源
     * 3.执行初始化脚本
     * 4.返回封装后的数据源
     * @param dataSourceProperty 数据源属性
     * @return
     */
    @Override
    public DataSource createDataSource(DataSourceProperty dataSourceProperty) {
        String publicKey = dataSourceProperty.getPublicKey();
        if (StringUtils.isEmpty(publicKey)) {
            publicKey = properties.getPublicKey();
            dataSourceProperty.setPublicKey(publicKey);
        }
        Boolean lazy  = dataSourceProperty.getLazy();
        if (lazy == null) {
            lazy = properties.getLazy();
            dataSourceProperty.setLazy(lazy);
        }
        dataSourceInitEvent.beforeCreate(dataSourceProperty);
        DataSource dataSource = doCreateDataSource(dataSourceProperty);
        dataSourceInitEvent.afterCreate(dataSource);
        this.runScrip(dataSource, dataSourceProperty);
        return wrapDataSource(dataSource, dataSourceProperty);
    }

    /**
     * 支持创建完成数据源连接之后,初始化脚本
     * @param dataSource
     * @param dataSourceProperty
     */
    private void runScrip(DataSource dataSource, DataSourceProperty dataSourceProperty) {
        DatasourceInitProperties initProperty = dataSourceProperty.getInit();
        String schema = initProperty.getSchema();
        String data = initProperty.getData();
        if (StringUtils.hasText(schema) || StringUtils.hasText(data)) {
            ScriptRunner scriptRunner = new ScriptRunner(initProperty.isContinueOnError(), initProperty.getSeparator());
            if (StringUtils.hasText(schema)) {
                scriptRunner.runScript(dataSource, schema);
            }
            if (StringUtils.hasText(data)) {
                scriptRunner.runScript(dataSource, data);
            }
        }
    }

    private DataSource wrapDataSource(DataSource dataSource, DataSourceProperty dataSourceProperty) {
        String name = dataSourceProperty.getPoolName();
        DataSource targetDataSource = dataSource;

        Boolean enabledP6spy = properties.getP6spy() && dataSourceProperty.getP6spy();
        if (enabledP6spy) {
            targetDataSource = new P6DataSource(dataSource);
            log.debug("dynamic-datasource [{}] wrap p6spy plugin", name);
        }

        Boolean enabledSeata = properties.getSeata() && dataSourceProperty.getSeata();
        SeataMode seataMode = properties.getSeataMode();
        if (enabledSeata) {
            if (SeataMode.XA == seataMode) {
                targetDataSource = new DataSourceProxyXA(targetDataSource);
            } else {
                targetDataSource = new DataSourceProxy(targetDataSource);
            }
            log.debug("dynamic-datasource [{}] wrap seata plugin transaction mode ", name);
        }
        return new ItemDataSource(name, dataSource, targetDataSource, enabledP6spy, enabledSeata, seataMode);
    }
}
复制代码

数据源提供者

数据源提供者是连接配置文件和数据源创建器的桥梁

/**
 * 多数据源加载接口,默认的实现为从yml信息中加载所有数据源 你可以自己实现从其他地方加载所有数据源
 *
 * @author TaoYu Kanyuxia
 * @see YmlDynamicDataSourceProvider
 * @see AbstractJdbcDataSourceProvider
 * @since 1.0.0
 */
public interface DynamicDataSourceProvider {

    /**
     * 加载所有数据源
     *
     * @return 所有数据源,key为数据源名称
     */
    Map<String, DataSource> loadDataSources();
}
复制代码
@Slf4j
public abstract class AbstractDataSourceProvider implements DynamicDataSourceProvider {

    /**
     * 数据源创建器
     */
    @Autowired
    private DefaultDataSourceCreator defaultDataSourceCreator;

    /**
     * 循环遍历从配置文件读取的多个数据源, 然后根据数据源的类型, 调用DataSourceCreator数据源创建器去创建(初始化)数据源,
     * 然后返回已经初始化好的数据源,将其保存到map集合中.
     * @param dataSourcePropertiesMap
     * @return
     */
    protected Map<String, DataSource> createDataSourceMap(
            Map<String, DataSourceProperty> dataSourcePropertiesMap) {
        Map<String, DataSource> dataSourceMap = new HashMap<>(dataSourcePropertiesMap.size() * 2);
        for (Map.Entry<String, DataSourceProperty> item : dataSourcePropertiesMap.entrySet()) {
            String dsName = item.getKey();
            DataSourceProperty dataSourceProperty = item.getValue();
            String poolName = dataSourceProperty.getPoolName();
            if (poolName == null || "".equals(poolName)) {
                poolName = dsName;
            }
            dataSourceProperty.setPoolName(poolName);
            dataSourceMap.put(dsName, defaultDataSourceCreator.createDataSource(dataSourceProperty));
        }
        return dataSourceMap;
    }
}
复制代码
/**
 * YML数据源提供者
 *
 * @author TaoYu Kanyuxia
 * @since 1.0.0
 */
@Slf4j
@AllArgsConstructor
public class YmlDynamicDataSourceProvider extends AbstractDataSourceProvider {

    /**
     * 所有数据源
     */
    private final Map<String, DataSourceProperty> dataSourcePropertiesMap;

    @Override
    public Map<String, DataSource> loadDataSources() {
        return createDataSourceMap(dataSourcePropertiesMap);
    }
}
复制代码

动态路由数据源

这里是动态切换数据源的核心.下面是类图结构

首先来看看AbstractRoutingDataSource抽象类

/**
 * 抽象动态获取数据源
 *
 * @author TaoYu
 * @since 2.2.0
 */
public abstract class AbstractRoutingDataSource extends AbstractDataSource {

    /**
     * 抽象获取连接池
     *
     * @return 连接池
     */
    protected abstract DataSource determineDataSource();

    /**
     * 获取数据库连接 如果连接不存在 则创建连接
     * 并维护连接与ds的关系 方便下次使用
     * @return
     * @throws SQLException
     */
    @Override
    public Connection getConnection() throws SQLException {
        String xid = TransactionContext.getXID();
        if (StringUtils.isEmpty(xid)) {
            return determineDataSource().getConnection();
        } else {
            String ds = DynamicDataSourceContextHolder.peek();
            ds = StringUtils.isEmpty(ds) ? "default" : ds;
            ConnectionProxy connection = ConnectionFactory.getConnection(ds);
            return connection == null ? getConnectionProxy(ds, determineDataSource().getConnection()) : connection;
        }
    }

    /**
     * 通过账号密码
     * 获取数据库连接 如果连接不存在 则创建连接
     * 并维护连接与ds的关系 方便下次使用
     * @return
     * @throws SQLException
     */
    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        String xid = TransactionContext.getXID();
        if (StringUtils.isEmpty(xid)) {
            return determineDataSource().getConnection(username, password);
        } else {
            String ds = DynamicDataSourceContextHolder.peek();
            ds = StringUtils.isEmpty(ds) ? "default" : ds;
            ConnectionProxy connection = ConnectionFactory.getConnection(ds);
            return connection == null ? getConnectionProxy(ds, determineDataSource().getConnection(username, password))
                    : connection;
        }
    }

    /**
     * 新建连接的代理对象
     * @param ds
     * @param connection
     * @return
     */
    private Connection getConnectionProxy(String ds, Connection connection) {
        ConnectionProxy connectionProxy = new ConnectionProxy(connection, ds);
        ConnectionFactory.putConnection(ds, connectionProxy);
        return connectionProxy;
    }

    @Override
    @SuppressWarnings("unchecked")
    public <T> T unwrap(Class<T> iface) throws SQLException {
        if (iface.isInstance(this)) {
            return (T) this;
        }
        return determineDataSource().unwrap(iface);
    }

    @Override
    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return (iface.isInstance(this) || determineDataSource().isWrapperFor(iface));
    }
}
复制代码

下面是DynamicRoutingDataSource,提供了对数据源进行管理的相关方法

public class DynamicRoutingDataSource extends AbstractRoutingDataSource implements InitializingBean, DisposableBean {
    /**
     * 数据源分组分隔符
     * eg:master_1,master_2
     */
    private static final String UNDERLINE = "_";
    /**
     * 所有数据库
     */
    private final Map<String, DataSource> dataSourceMap = new ConcurrentHashMap<>();
    /**
     * 分组数据库
     */
    private final Map<String, GroupDataSource> groupDataSources = new ConcurrentHashMap<>();
    /**
     * 数据源提供者集合 在afterPropertiesSet()方法执行时
     * 项目启动时加载数据源
     */
    @Autowired
    private List<DynamicDataSourceProvider> providers;
    /**
     * 数据源组中的数据源防伪策略
     * 默认采用负载均衡策略
     * 框架中还实现了随机策略
     */
    @Setter
    private Class<? extends DynamicDataSourceStrategy> strategy = LoadBalanceDynamicDataSourceStrategy.class;
    /**
     * 默认数据源的dsName
     */
    @Setter
    private String primary = "master";
    /**
     * 是否启用严格模式,默认不启动. 严格模式下未匹配到数据源直接报错, 非严格模式下则使用默认数据源primary所设置的数据源
     */
    @Setter
    private Boolean strict = false;
    /**
     * 是否使用p6spy输出,默认不输出
     */
    @Setter
    private Boolean p6spy = false;
    /**
     * 是否使用开启seata,默认不开启
     */
    @Setter
    private Boolean seata = false;

    /**
     * 决定数据源 从ThreadLocal中获取数据源的名称
     * 返回对应的数据源 在方法上层会根据数据源获取对应的连接
     * @return
     */
    @Override
    public DataSource determineDataSource() {
        String dsKey = DynamicDataSourceContextHolder.peek();
        return getDataSource(dsKey);
    }

    /**
     * 确定使用默认数据源
     * @return
     */
    private DataSource determinePrimaryDataSource() {
        log.debug("dynamic-datasource switch to the primary datasource");
        DataSource dataSource = dataSourceMap.get(primary);
        if (dataSource != null) {
            return dataSource;
        }
        GroupDataSource groupDataSource = groupDataSources.get(primary);
        if (groupDataSource != null) {
            return groupDataSource.determineDataSource();
        }
        throw new CannotFindDataSourceException("dynamic-datasource can not find primary datasource");
    }

    /**
     * 获取所有的数据源
     *
     * @return 当前所有数据源
     */
    public Map<String, DataSource> getDataSources() {
        return dataSourceMap;
    }

    /**
     * 获取的所有的分组数据源
     *
     * @return 当前所有的分组数据源
     */
    public Map<String, GroupDataSource> getGroupDataSources() {
        return groupDataSources;
    }

    /**
     * 获取数据源
     *
     * @param ds 数据源名称
     * @return 数据源
     */
    public DataSource getDataSource(String ds) {
        if (StringUtils.isEmpty(ds)) {
            return determinePrimaryDataSource();
        } else if (!groupDataSources.isEmpty() && groupDataSources.containsKey(ds)) {
            log.debug("dynamic-datasource switch to the datasource named [{}]", ds);
            return groupDataSources.get(ds).determineDataSource();
        } else if (dataSourceMap.containsKey(ds)) {
            log.debug("dynamic-datasource switch to the datasource named [{}]", ds);
            return dataSourceMap.get(ds);
        }
        if (strict) {
            throw new CannotFindDataSourceException("dynamic-datasource could not find a datasource named" + ds);
        }
        return determinePrimaryDataSource();
    }

    /**
     * 添加数据源
     *
     * @param ds         数据源名称
     * @param dataSource 数据源
     */
    public synchronized void addDataSource(String ds, DataSource dataSource) {
        DataSource oldDataSource = dataSourceMap.put(ds, dataSource);
        // 新数据源添加到分组
        this.addGroupDataSource(ds, dataSource);
        // 关闭老的数据源
        if (oldDataSource != null) {
            closeDataSource(ds, oldDataSource);
        }
        log.info("dynamic-datasource - add a datasource named [{}] success", ds);
    }

    /**
     * 新数据源添加到分组
     *
     * @param ds         新数据源的名字
     * @param dataSource 新数据源
     */
    private void addGroupDataSource(String ds, DataSource dataSource) {
        if (ds.contains(UNDERLINE)) {
            String group = ds.split(UNDERLINE)[0];
            GroupDataSource groupDataSource = groupDataSources.get(group);
            if (groupDataSource == null) {
                try {
                    groupDataSource = new GroupDataSource(group, strategy.getDeclaredConstructor().newInstance());
                    groupDataSources.put(group, groupDataSource);
                } catch (Exception e) {
                    throw new RuntimeException("dynamic-datasource - add the datasource named " + ds + " error", e);
                }
            }
            groupDataSource.addDatasource(ds, dataSource);
        }
    }

    /**
     * 删除数据源
     *
     * @param ds 数据源名称
     */
    public synchronized void removeDataSource(String ds) {
        if (!StringUtils.hasText(ds)) {
            throw new RuntimeException("remove parameter could not be empty");
        }
        if (primary.equals(ds)) {
            throw new RuntimeException("could not remove primary datasource");
        }
        if (dataSourceMap.containsKey(ds)) {
            DataSource dataSource = dataSourceMap.remove(ds);
            closeDataSource(ds, dataSource);
            if (ds.contains(UNDERLINE)) {
                String group = ds.split(UNDERLINE)[0];
                if (groupDataSources.containsKey(group)) {
                    DataSource oldDataSource = groupDataSources.get(group).removeDatasource(ds);
                    if (oldDataSource == null) {
                        log.warn("fail for remove datasource from group. dataSource: {} ,group: {}", ds, group);
                    }
                }
            }
            log.info("dynamic-datasource - remove the database named [{}] success", ds);
        } else {
            log.warn("dynamic-datasource - could not find a database named [{}]", ds);
        }
    }

    @Override
    public void destroy() throws Exception {
        log.info("dynamic-datasource start closing ....");
        for (Map.Entry<String, DataSource> item : dataSourceMap.entrySet()) {
            closeDataSource(item.getKey(), item.getValue());
        }
        log.info("dynamic-datasource all closed success,bye");
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        // 检查开启了配置但没有相关依赖
        checkEnv();
        // 添加并分组数据源
        Map<String, DataSource> dataSources = new HashMap<>(16);
        for (DynamicDataSourceProvider provider : providers) {
            dataSources.putAll(provider.loadDataSources());
        }
        for (Map.Entry<String, DataSource> dsItem : dataSources.entrySet()) {
            addDataSource(dsItem.getKey(), dsItem.getValue());
        }
        // 检测默认数据源是否设置
        if (groupDataSources.containsKey(primary)) {
            log.info("dynamic-datasource initial loaded [{}] datasource,primary group datasource named [{}]", dataSources.size(), primary);
        } else if (dataSourceMap.containsKey(primary)) {
            log.info("dynamic-datasource initial loaded [{}] datasource,primary datasource named [{}]", dataSources.size(), primary);
        } else {
            log.warn("dynamic-datasource initial loaded [{}] datasource,Please add your primary datasource or check your configuration", dataSources.size());
        }
    }

    private void checkEnv() {
        if (p6spy) {
            try {
                Class.forName("com.p6spy.engine.spy.P6DataSource");
                log.info("dynamic-datasource detect P6SPY plugin and enabled it");
            } catch (Exception e) {
                throw new RuntimeException("dynamic-datasource enabled P6SPY ,however without p6spy dependency", e);
            }
        }
        if (seata) {
            try {
                Class.forName("io.seata.rm.datasource.DataSourceProxy");
                log.info("dynamic-datasource detect ALIBABA SEATA and enabled it");
            } catch (Exception e) {
                throw new RuntimeException("dynamic-datasource enabled ALIBABA SEATA,however without seata dependency", e);
            }
        }
    }

    /**
     * close db
     *
     * @param ds         dsName
     * @param dataSource db
     */
    private void closeDataSource(String ds, DataSource dataSource) {
        try {
            if (dataSource instanceof ItemDataSource) {
                ((ItemDataSource) dataSource).close();
            } else {
                if (seata) {
                    if (dataSource instanceof DataSourceProxy) {
                        DataSourceProxy dataSourceProxy = (DataSourceProxy) dataSource;
                        dataSource = dataSourceProxy.getTargetDataSource();
                    }
                }
                if (p6spy) {
                    if (dataSource instanceof P6DataSource) {
                        Field realDataSourceField = P6DataSource.class.getDeclaredField("realDataSource");
                        realDataSourceField.setAccessible(true);
                        dataSource = (DataSource) realDataSourceField.get(dataSource);
                    }
                }
                Method closeMethod = ReflectionUtils.findMethod(dataSource.getClass(), "close");
                if (closeMethod != null) {
                    closeMethod.invoke(dataSource);
                }
            }
        } catch (Exception e) {
            log.warn("dynamic-datasource closed datasource named [{}] failed", ds, e);
        }
    }

}
复制代码
/**
 * 核心基于ThreadLocal的切换数据源工具类
 *
 * @author TaoYu Kanyuxia
 * @since 1.0.0
 */
public final class DynamicDataSourceContextHolder {

    /**
     * 为什么要用链表存储(准确的是栈)
     * <pre>
     * 为了支持嵌套切换,如ABC三个service都是不同的数据源
     * 其中A的某个业务要调B的方法,B的方法需要调用C的方法。一级一级调用切换,形成了链。
     * 传统的只设置当前线程的方式不能满足此业务需求,必须使用栈,后进先出。
     * </pre>
     */
    private static final ThreadLocal<Deque<String>> LOOKUP_KEY_HOLDER = new NamedThreadLocal<Deque<String>>("dynamic-datasource") {
        @Override
        protected Deque<String> initialValue() {
            return new ArrayDeque<>();
        }
    };

    private DynamicDataSourceContextHolder() {
    }

    /**
     * 获得当前线程数据源
     *
     * @return 数据源名称
     */
    public static String peek() {
        return LOOKUP_KEY_HOLDER.get().peek();
    }

    /**
     * 设置当前线程数据源
     * <p>
     * 如非必要不要手动调用,调用后确保最终清除
     * </p>
     *
     * @param ds 数据源名称
     */
    public static String push(String ds) {
        String dataSourceStr = StringUtils.isEmpty(ds) ? "" : ds;
        LOOKUP_KEY_HOLDER.get().push(dataSourceStr);
        return dataSourceStr;
    }

    /**
     * 清空当前线程数据源
     * <p>
     * 如果当前线程是连续切换数据源 只会移除掉当前线程的数据源名称
     * </p>
     */
    public static void poll() {
        Deque<String> deque = LOOKUP_KEY_HOLDER.get();
        deque.poll();
        if (deque.isEmpty()) {
            LOOKUP_KEY_HOLDER.remove();
        }
    }

    /**
     * 强制清空本地线程
     * <p>
     * 防止内存泄漏,如手动调用了push可调用此方法确保清除
     * </p>
     */
    public static void clear() {
        LOOKUP_KEY_HOLDER.remove();
    }
}
复制代码

猜你喜欢

转载自juejin.im/post/7096744699259519012