Spring 动态切换、添加数据源实现以及源码浅析

程序员小乐 2020-10-08

点击上方 "程序员小乐" ，关注公众号

8点20分，第一时间与你相约

每日英文

There are things that we don't want to happen but have to accept,things we don't want to know but have to learn,and people we can't live without but have to let go.

总有一些事,我们不愿它发生,却必须接受；总有些东西,我们不想知道, 却必须了解；总有些人, 我们不能没有, 却必须学着放手。

每日掏心话

人生之所以精彩，是他愿意全然的接受一切。生命之所以可贵，是他愿意尊重一切的生命。

来自：殷天文 | 责编：乐乐

链接：jianshu.com/p/0a485c965b8b

图片来自网络

往日回顾：为什么？为什么程序员总是加班！

正文

公司项目需求，由于要兼容老系统的数据库结构，需要搭建一个可以动态切换、添加数据源的后端服务。

参考了过去的项目，通过配置多个SqlSessionFactory 来实现多数据源，这么做的话，未免过于笨重，而且无法实现动态添加数据源这个需求。

通过 spring AbstractRoutingDataSource 为我们抽象了一个 DynamicDataSource 解决这一问题

简单分析下 AbstractRoutingDataSource 的源码

targetDataSources 就是我们的多个数据源，在初始化的时候会调用afterPropertiesSet()，去解析我们的数据源然后 put 到 resolvedDataSources

实现了 DataSource 的 getConnection(); 我们看看 determineTargetDataSource(); 做了什么

通过下面的 determineCurrentLookupKey();（这个方法需要我们实现）返回一个key，然后从 resolvedDataSources （其实也就是 targetDataSources）中 get 一个数据源，实现了每次调用 getConnection(); 打开连接切换数据源，如果想动态添加的话只需要重新 set targetDataSources 再调用 afterPropertiesSet() 即可

Talk is cheap. Show me the code

我使用的springboot版本为 1.5.x，下面是核心代码

完整代码：

https://gitee.com/yintianwen7/spring-dynamic-datasource

/**
 * 多数据源配置
 * 
 * @author Taven
 *
 */
@Configuration
@MapperScan("com.gitee.taven.mapper")
public class DataSourceConfigurer {

    /**
     * DataSource 自动配置并注册
     *
     * @return data source
     */
    @Bean("db0")
    @Primary
    @ConfigurationProperties(prefix = "datasource.db0")
    public DataSource dataSource0() {
        return DruidDataSourceBuilder.create().build();
    }

    /**
     * DataSource 自动配置并注册
     *
     * @return data source
     */
    @Bean("db1")
    @ConfigurationProperties(prefix = "datasource.db1")
    public DataSource dataSource1() {
        return DruidDataSourceBuilder.create().build();
    }

    /**
     * 注册动态数据源
     * 
     * @return
     */
    @Bean("dynamicDataSource")
    public DataSource dynamicDataSource() {
        DynamicRoutingDataSource dynamicRoutingDataSource = new DynamicRoutingDataSource();
        Map<Object, Object> dataSourceMap = new HashMap<>();
        dataSourceMap.put("dynamic_db0", dataSource0());
        dataSourceMap.put("dynamic_db1", dataSource1());
        dynamicRoutingDataSource.setDefaultTargetDataSource(dataSource0());// 设置默认数据源
        dynamicRoutingDataSource.setTargetDataSources(dataSourceMap);
        return dynamicRoutingDataSource;
    }

    /**
     * Sql session factory bean.
     * Here to config datasource for SqlSessionFactory
     * <p>
     * You need to add @{@code @ConfigurationProperties(prefix = "mybatis")}, if you are using *.xml file,
     * the {@code 'mybatis.type-aliases-package'} and {@code 'mybatis.mapper-locations'} should be set in
     * {@code 'application.properties'} file, or there will appear invalid bond statement exception
     *
     * @return the sql session factory bean
     */
    @Bean
    @ConfigurationProperties(prefix = "mybatis")
    public SqlSessionFactoryBean sqlSessionFactoryBean() {
        SqlSessionFactoryBean sqlSessionFactoryBean = new SqlSessionFactoryBean();
        // 必须将动态数据源添加到 sqlSessionFactoryBean
        sqlSessionFactoryBean.setDataSource(dynamicDataSource());
        return sqlSessionFactoryBean;
    }

    /**
     * 事务管理器
     *
     * @return the platform transaction manager
     */
    @Bean
    public PlatformTransactionManager transactionManager() {
        return new DataSourceTransactionManager(dynamicDataSource());
    }
}

通过 ThreadLocal 获取线程安全的数据源 key

package com.gitee.taven.config;

public class DynamicDataSourceContextHolder {

    private static final ThreadLocal<String> contextHolder = new ThreadLocal<String>() {
        @Override
        protected String initialValue() {
            return "dynamic_db0";
        }
    };

    /**
     * To switch DataSource
     *
     * @param key the key
     */
    public static void setDataSourceKey(String key) {
        contextHolder.set(key);
    }

    /**
     * Get current DataSource
     *
     * @return data source key
     */
    public static String getDataSourceKey() {
        return contextHolder.get();
    }

    /**
     * To set DataSource as default
     */
    public static void clearDataSourceKey() {
        contextHolder.remove();
    }

}

动态添加、切换数据源

/**
 * 动态数据源
 * 
 * @author Taven
 *
 */
public class DynamicRoutingDataSource extends AbstractRoutingDataSource {

    private final Logger logger = LoggerFactory.getLogger(getClass());

    private static Map<Object, Object> targetDataSources = new HashMap<>();

    /**
     * 设置当前数据源
     *
     * @return
     */
    @Override
    protected Object determineCurrentLookupKey() {
        logger.info("Current DataSource is [{}]", DynamicDataSourceContextHolder.getDataSourceKey());
        return DynamicDataSourceContextHolder.getDataSourceKey();
    }

    @Override
    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        super.setTargetDataSources(targetDataSources);
        DynamicRoutingDataSource.targetDataSources = targetDataSources;
    }

    /**
     * 是否存在当前key的 DataSource
     * 
     * @param key
     * @return 存在返回 true, 不存在返回 false
     */
    public static boolean isExistDataSource(String key) {
        return targetDataSources.containsKey(key);
    }

    /**
     * 动态增加数据源
     * 
     * @param map 数据源属性
     * @return
     */
    public synchronized boolean addDataSource(Map<String, String> map) {
        try {
            Connection connection = null;
            // 排除连接不上的错误
            try { 
                Class.forName(map.get(DruidDataSourceFactory.PROP_DRIVERCLASSNAME));
                connection = DriverManager.getConnection(
                        map.get(DruidDataSourceFactory.PROP_URL), 
                        map.get(DruidDataSourceFactory.PROP_USERNAME),
                        map.get(DruidDataSourceFactory.PROP_PASSWORD));
                System.out.println(connection.isClosed());
            } catch (Exception e) {
                return false;
            } finally {
                if (connection != null && !connection.isClosed()) 
                    connection.close();
            }
            String database = map.get("database");//获取要添加的数据库名
            if (StringUtils.isBlank(database)) return false;
            if (DynamicRoutingDataSource.isExistDataSource(database)) return true; 
            DruidDataSource druidDataSource = (DruidDataSource) DruidDataSourceFactory.createDataSource(map);
            druidDataSource.init();
            Map<Object, Object> targetMap = DynamicRoutingDataSource.targetDataSources;
            targetMap.put(database, druidDataSource);
            // 当前 targetDataSources 与 父类 targetDataSources 为同一对象 所以不需要set
//          this.setTargetDataSources(targetMap);
            this.afterPropertiesSet();
            logger.info("dataSource {} has been added", database);
        } catch (Exception e) {
            logger.error(e.getMessage());
            return false;
        }
        return true;
    }

}