AbstractRoutingDataSource 实现动态数据源切换原理简单分析

AbstractRoutingDataSource 实现动态数据源切换原理简单分析

写在前面，项目中用到了动态数据源切换，记录一下其运行机制。

代码展示

下面列出一些关键代码，后续分析会用到

1. 数据配置

@Configuration
@PropertySource({ "classpath:jdbc.yml" })
@EnableTransactionManagement(proxyTargetClass = true)
public class DataConfig {

    @Autowired
    private Environment env ;

    /**
     *  将jdbc相关的异常转换为spring的异常类型
     */
    @Bean
    public BeanPostProcessor persistenceTransLation(){
        return new PersistenceExceptionTranslationPostProcessor() ;
    }

    /**
     * 多数据源
     * @return
     */
    @Bean
    public DynamicDataSource dynamicDataSource(){
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        Map<Object,Object> sourceMap = new HashMap<>();
        //取得所有的datasource，DataSourceEnum里存放数据源的唯一标识
        EnumSet<DataSourceEnum> enums = EnumSet.allOf(DataSourceEnum.class);
        for(DataSourceEnum dataSource:enums){
            // map存放数据源的key和数据源
            sourceMap.put(dataSource.getKey(),generateDataSource(dataSource.getKey()));
        }
        //  重点
        dynamicDataSource.setTargetDataSources(sourceMap);
     dynamicDataSource.setDefaultTargetDataSource(sourceMap.get(DataSourceEnum.TEST.getKey()));
        return dynamicDataSource;
    }

    // 读取配置文件，创建数据源对象
    private EncryptDataSource generateDataSource(String key){
        EncryptDataSource dataSource
                = new EncryptDataSource();
        key = key.toLowerCase() ;
        String url = "jdbc.url."+key;
        String username = "jdbc.username."+key;
        String password = "jdbc.password."+key;
        dataSource.setDriverClassName("com.sybase.jdbc4.jdbc.SybDataSource");//SybDriver
        dataSource.setUrl(env.getProperty(url));
        dataSource.setUsername(env.getProperty(username));
        dataSource.setPassword(env.getProperty(password));

        //配置连接池
        dataSource.setInitialSize(Integer.parseInt(env.getProperty("jdbc.initialSize")));
        dataSource.setMaxIdle(Integer.parseInt(env.getProperty("jdbc.maxIdle")));
        dataSource.setMinIdle(Integer.parseInt(env.getProperty("jdbc.minIdle")));
        return dataSource;
    }
}

1. 自定义数据源类

public class DynamicDataSource extends AbstractRoutingDataSource {
    // 存放数据源的id（唯一标识）
    private static final ThreadLocal<String> dataSourceHolder = new ThreadLocal<>() ;

    //  重点
    @Override
    protected Object determineCurrentLookupKey() {
        return dataSourceHolder.get();
    }

    // 切换数据源
    public static void router(String sourceKey){
        if(StrUtil.isEmpty(sourceKey)){
            return;
        }
        if(DataSourceEnum.getSourceByKey(sourceKey)!=null){
            //根据法院代码切换
            dataSourceHolder.set(DataSourceEnum.getSourceByKey(sourceKey));
        }
    }

    ……

}

1. 数据源配置（jdbc.yml）

#测试库
jdbc.url.test: jdbc:sybase:Tds:xxx.xxx.xxx.xxx:xx/JUDGE?charset=cp936
jdbc.username.test: fymis
jdbc.password.test: xx

原理分析

第一部分已将关键代码列出，该部分通过修改后即可实现数据源的切换功能。下面来分析一下流程。

AbstractRoutingDataSource 类解析

只列出了部分方法，需要详细代码请自行移步源码

public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {
    @Nullable
    private Map<Object, Object> targetDataSources;// 目标数据源map
    @Nullable
    private Object defaultTargetDataSource;// 默认数据源
    private boolean lenientFallback = true;
    private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();
    @Nullable
    private Map<Object, DataSource> resolvedDataSources;
    @Nullable
    private DataSource resolvedDefaultDataSource;

    public AbstractRoutingDataSource() {
    }

    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        this.targetDataSources = targetDataSources;
    }

    public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
        this.defaultTargetDataSource = defaultTargetDataSource;
    }

    // 初始化 Bean 时执行
    public void afterPropertiesSet() {
        if (this.targetDataSources == null) {
            throw new IllegalArgumentException("Property 'targetDataSources' is required");
        } else {
            // 将targetDataSources属性的值赋值给resolvedDataSources，后续需要用到resolvedDataSources
            this.resolvedDataSources = new HashMap(this.targetDataSources.size());
            this.targetDataSources.forEach((key, value) -> {
                Object lookupKey = this.resolveSpecifiedLookupKey(key);
                DataSource dataSource = this.resolveSpecifiedDataSource(value);
                this.resolvedDataSources.put(lookupKey, dataSource);
            });
            if (this.defaultTargetDataSource != null) {
                this.resolvedDefaultDataSource = this.resolveSpecifiedDataSource(this.defaultTargetDataSource);
            }

        }
    }

    // 重写了 getConnection 方法，ORM 框架执行语句前会调用该处
    @Override
    public Connection getConnection() throws SQLException {
        return this.determineTargetDataSource().getConnection();
    }

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

    //  重点
    protected DataSource determineTargetDataSource() {
        Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
        // 调用我们重写的determineCurrentLookupKey方法，返回的是数据源的唯一标识
        Object lookupKey = this.determineCurrentLookupKey();
        // 从map中查询改标识对应的数据源，然后返回该数据源，打开对应连接
        DataSource 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 + "]");
        } else {
            return dataSource;
        }
    }

    // 钩子方法，供我们重写
    @Nullable
    protected abstract Object determineCurrentLookupKey();
}

总结与闲谈

综上，可以列出以下几点描述整个流程：

1. 自定义类继承 AbstractRoutingDataSource（后文称 ARDS），重写 determineCurrentLookupKey()，返回数据源的唯一标识；
2. 将数据源名称和数据源封装为 map，调用 ARDS 类的 setTargetDataSources() 设置目标数据源。ARDS 类实现了 InitializingBean 接口，重写了 afterPropertySet()（对该方法不熟悉的话请回顾一下 Bean 的生命周期，该方法在 Bean 的属性注入后执行），该方法内部对 resolvedDataSources 属性赋值（将 targetDataSources 的值放进去），后续会用到 resolvedDataSources ；
3. ARDS 实现了 DataSource 接口，重写了 getConnection()，当 ORM 框架执行 sql 语句前总是执行 getConnection()，然后就调用到了重写后的 getConnection()，该方法内部调用了 ARDS 类的 determineTargetDataSource()；
4. determineTargetDataSource() 内部调用了自定义类重写的 determineCurrentLookupKey()，返回数据源的映射，然后从 resolvedDataSources（map） 属性获取到数据源，进行后续的操作。

（题外话）想要实现数据源切换可以有两种实现：

1. 手动切换数据源，每次执行相应操作前调用 router 方法切换；
2. 还有一种思路就是利用 AOP，设计一个注解，注解内添加数据源唯一标识的属性，然后对方法添加注解，AOP 代码进行拦截，然后将唯一标识赋值给 ThreadLocal 变量即可。