Spring代码中动态切换数据源

最近项目中遇到一个场景，需要能够在一个方法中操作多个具有相同表结构数据库(你可以理解为一个表中的数据被水平拆分到多个库中，查询时需要遍历这多个库)。经过笔者几天的研究，最终解决了问题，并且写了一个demo共享到我的github。

关注笔者博客的小伙伴一定知道之前的这篇文章点击打开链接，这篇博客中的解决方案仅仅适用读写分离的场景。就是说，当你在开发的时候已经确定了使用写库一读库的形式。笔者今天要写的这篇文章具有普适性，适合所有需要在Spring工程中动态切换数据源的场景，而且本文中的解决方案对工程的代码基本没有侵入性。下面就来说下该方案的实现原理：

在Spring-Mybatis中，有这样一个类AbstractRoutingDataSource，根据名字可以猜到，这是一个框架提供的用于动态选择数据源的类。这个类有两个重要的参数，分别叫

defaultTargetDataSource和targetDataSources。一般的工程都是一个数据源，所以不太接触到这个类。在作者之前的博客自动切换多个数据源中，可以看到这个类的xml配置如下：

    <bean id="myoneDataSource" class="org.apache.commons.dbcp2.BasicDataSource" destroy-method="close">          <property name="driverClassName" value="${jdbc.myone.driver}"/>          <property name="url" value="${jdbc.myone.url}"/>          <property name="username" value="${jdbc.myone.username}"/>          <property name="password" value="${jdbc.myone.password}"/>      </bean>      <bean id="mytwoDataSource" class="org.apache.commons.dbcp2.BasicDataSource" destroy-method="close">          <property name="driverClassName" value="${jdbc.mytwo.driver}"/>          <property name="url" value="${jdbc.mytwo.url}"/>          <property name="username" value="${jdbc.mytwo.username}"/>          <property name="password" value="${jdbc.mytwo.password}"/>      </bean>        <bean id="multipleDataSource" class="dal.datasourceswitch.MultipleDataSource">          <property name="defaultTargetDataSource" ref="myoneDataSource"/> <!--默认主库-->          <property name="targetDataSources">              <map>                  <entry key="myone" value-ref="myoneDataSource"/>            <!--辅助aop完成自动数据库切换-->                  <entry key="mytwo" value-ref="mytwoDataSource"/>              </map>          </property>      </bean> 

上面的配置文件对这两个参数的描述已经很清楚了，但这是多个数据源已经确定的场景。我们这篇博客中的场景是多个数据源的信息存在于数据库中，可能数据库中的数据源信息会动态的增加或者减少。这样的话，就不能像上面这样配置了。那怎么办呢？

我们仅仅需要设定默认的数据源，即defaultDataSource参数，至于targetDataSources参数我们需要在代码中动态的设定。来看下具体的xml配置：

<bean id="defaultDataSource" class="org.apache.commons.dbcp.BasicDataSource" destroy-method="close"          p:driverClassName="${db_driver}"          p:url="${db_url}"          p:username="${db_user}"          p:password="${db_pass}"          p:validationQuery="select 1"          p:testOnBorrow="true"/>    <!--动态数据源相关-->    <bean id="dynamicDataSource" class="org.xyz.test.service.datasourceswitch.impl.DynamicDataSource">        <property name="targetDataSources">            <map key-type="java.lang.String">                <entry key="defaultDataSource" value-ref="defaultDataSource"/>            </map>        </property>        <property name="defaultTargetDataSource" ref="defaultDataSource"/>    </bean>

从上面的配置文件中可以看到，我们仅仅配置了默认的数据源defaultDataSource。至于其他的数据源targetDataSources，我们没有配置，需要在代码中动态的创建。关于配置就讲清楚啦！但我们注意到，支持动态数据源的不应该是AbstractRoutingDataSource类吗？怎么上面的配置中是DynamicDataSource类。没错，这个是我们自定义的继承自AbstractRoutingDataSource类的类，也只最重要的类，来看下：（理解这个类，你需要熟练掌握JAVA反射，以及ThreadLocal变量，和Spring的注入机制。别退缩，大家都是这样一步步学过来的！）（下面仅仅是看下全貌，代码的下面会有详细的说明）

final class DynamicDataSource extends AbstractRoutingDataSource implements ApplicationContextAware{    private static final String DATA_SOURCES_NAME = "targetDataSources";    private ApplicationContext applicationContext;    @Override    protected Object determineCurrentLookupKey() {        DataSourceBeanBuilder dataSourceBeanBuilder = DataSourceHolder.getDataSource();        System.out.println("----determineCurrentLookupKey---"+dataSourceBeanBuilder);        if (dataSourceBeanBuilder == null) {            return null;        }        DataSourceBean dataSourceBean = new DataSourceBean(dataSourceBeanBuilder);        //查看当前容器中是否存在        try {            Map<Object,Object> map=getTargetDataSources();            synchronized (this) {                if (!map.keySet().contains(dataSourceBean.getBeanName())) {                    map.put(dataSourceBean.getBeanName(), createDataSource(dataSourceBean));                    super.afterPropertiesSet();//通知spring有bean更新                }            }            return dataSourceBean.getBeanName();        } catch (NoSuchFieldException | IllegalAccessException e) {            throw new SystemException(ErrorEnum.MULTI_DATASOURCE_SWITCH_EXCEPTION);        }    }    private Object createDataSource(DataSourceBean dataSourceBean) throws IllegalAccessException {        //在spring容器中创建并且声明bean        ConfigurableApplicationContext context = (ConfigurableApplicationContext) applicationContext;        DefaultListableBeanFactory beanFactory = (DefaultListableBeanFactory) context.getBeanFactory();        BeanDefinitionBuilder beanDefinitionBuilder = BeanDefinitionBuilder.genericBeanDefinition(BasicDataSource.class);        //将dataSourceBean中的属性值赋给目标bean        Map<String, Object> properties = getPropertyKeyValues(DataSourceBean.class, dataSourceBean);        for (Map.Entry<String, Object> entry : properties.entrySet()) {            beanDefinitionBuilder.addPropertyValue((String) entry.getKey(), entry.getValue());        }        beanFactory.registerBeanDefinition(dataSourceBean.getBeanName(), beanDefinitionBuilder.getBeanDefinition());        return applicationContext.getBean(dataSourceBean.getBeanName());    }    private Map<Object, Object> getTargetDataSources() throws NoSuchFieldException, IllegalAccessException {        Field field = AbstractRoutingDataSource.class.getDeclaredField(DATA_SOURCES_NAME);        field.setAccessible(true);        return (Map<Object, Object>) field.get(this);    }    private <T> Map<String, Object> getPropertyKeyValues(Class<T> clazz, Object object) throws IllegalAccessException {        Field[] fields = clazz.getDeclaredFields();        Map<String, Object> result = new HashMap<>();        for (Field field : fields) {            field.setAccessible(true);            result.put(field.getName(), field.get(object));        }        result.remove("beanName");        return result;    }    @Override    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {        this.applicationContext=applicationContext;    }}

首先来看覆盖方法determineCurrentLookupKey()，框架在每次调用数据源时会先调用这个方法，以便知道使用哪个数据源。在本文的场景中，数据源是由程序员在即将切换数据源之前，将要使用的那个数据源的名称放到当前线程的ThreadLocal中，这样在determineCurrentLookupKey()方法中就可以从ThreadLocal中拿到当前请求钥匙用的数据源，从而进行初始化数据源并返回该数据源的操作。在ThreadLocal变量中，我们保存了一个DataSourceBuilder，这是一个建造者模式，不是本文的关键。我们在后面说。读者直接把他理解为是一个数据源的描述就好。因此，determineCurrentLookupKey()方法的流程就是：先从ThreadLocal中拿出要使用的数据源信息，然后看当前的targetDataSources中是否有了这个数据源。如果有直接返回。如果没有，创建一个这样的数据源，放到targetDataSources中然后返回。这个过程需要加锁，为何？这是典型的判断后插入场景，在多线程中会有线程安全问题，所以要加锁！至嘱！

由于targetDataSources是父类AbstractRoutingDataSource中的一个私有域，因此想要获得他的实例只能通过Java反射机制。这也是下面的方法存在的意义！

   private Map<Object, Object> getTargetDataSources() throws NoSuchFieldException, IllegalAccessException {        Field field = AbstractRoutingDataSource.class.getDeclaredField(DATA_SOURCES_NAME);        field.setAccessible(true);        return (Map<Object, Object>) field.get(this);    }

然后，我们来看具体是怎么创建数据源的。

 private Object createDataSource(DataSourceBean dataSourceBean) throws IllegalAccessException {        //在spring容器中创建并且声明bean        ConfigurableApplicationContext context = (ConfigurableApplicationContext) applicationContext;        DefaultListableBeanFactory beanFactory = (DefaultListableBeanFactory) context.getBeanFactory();        BeanDefinitionBuilder beanDefinitionBuilder = BeanDefinitionBuilder.genericBeanDefinition(BasicDataSource.class);        //将dataSourceBean中的属性值赋给目标bean        Map<String, Object> properties = getPropertyKeyValues(DataSourceBean.class, dataSourceBean);        for (Map.Entry<String, Object> entry : properties.entrySet()) {            beanDefinitionBuilder.addPropertyValue((String) entry.getKey(), entry.getValue());        }        beanFactory.registerBeanDefinition(dataSourceBean.getBeanName(), beanDefinitionBuilder.getBeanDefinition());        return applicationContext.getBean(dataSourceBean.getBeanName());    }

大家知道，Spring最主要的功能是作为bean容器，即他负责bean生命周期的管理。因此，我们自定义的datasource也不能“逍遥法外”，必须交给Spring容器来管理。这也正是DynamicDataSource类需要实现ApplicationContextAware并且注入ApplicationContext的原因。上面的代码就是根据指定的信息创建一个数据源。这种创建是Spring容器级别的创建。创建完毕之后，需要把刚刚创建的这个数据源放到targetDataSources中，并且还要通知Spring容器，targetDataSources对象变了。下面的方法就是在做这样的事情：

  private void addNewDataSourceToTargerDataSources(DataSourceBean dataSourceBean) throws NoSuchFieldException, IllegalAccessException {        getTargetDataSources().put(dataSourceBean.getBeanName(), createDataSource(dataSourceBean));        super.afterPropertiesSet();//通知spring有bean更新    }

上面的这一步很重要。没有这一步的话，Spring压根就不会知道targetDataSources中多了一个数据源。至此DynamicDataSource类就讲完了。其实仔细想想，思路还是很清晰的。啃掉了DynamicDataSource类这块硬骨头，下面就是一些辅助类了。比如说DataSourceHolder，业务代码通过使用这个类来通知DynamicDataSource中的determineCurrentLookupKey()方法到底使用那个数据源：

public final class DataSourceHolder {    private static ThreadLocal<DataSourceBeanBuilder> threadLocal=new ThreadLocal<DataSourceBeanBuilder>(){        @Override        protected DataSourceBeanBuilder initialValue() {            return null;        }    };    static DataSourceBeanBuilder getDataSource(){        return threadLocal.get();    }    public static void setDataSource(DataSourceBeanBuilder dataSourceBeanBuilder){        threadLocal.set(dataSourceBeanBuilder);    }    public static void clearDataSource(){        threadLocal.remove();    }}

再比如这个DataSourceBean类，实际上就是用于保存从那个默认数据库中拿出来的数据源信息，只不过为了安全起见，使用了builder模式，关于builder模式，可参见构建器模式。

final class DataSourceBean {    private final String beanName;    private final String driverClassName;    private final String url;    private final String username;    private final String password;    private final String validationQuery;    private final boolean testOnBorrow;    public DataSourceBean(DataSourceBeanBuilder beanBuilder){        this.beanName=beanBuilder.getBeanName();        this.driverClassName=beanBuilder.getDriverClassName();        this.url=beanBuilder.getUrl();        this.password=beanBuilder.getPassword();        this.testOnBorrow=beanBuilder.isTestOnBorrow();        this.username=beanBuilder.getUsername();        this.validationQuery=beanBuilder.getValidationQuery();    }    public String getBeanName() {        return beanName;    }    public String getDriverClassName() {        return driverClassName;    }    public String getUrl() {        return url;    }    public String getUsername() {        return username;    }    public String getPassword() {        return password;    }    public String getValidationQuery() {        return validationQuery;    }    public boolean isTestOnBorrow() {        return testOnBorrow;    }    @Override    public String toString() {        return "DataSourceBean{" +                "driverClassName='" + driverClassName + '\'' +                ", url='" + url + '\'' +                ", username='" + username + '\'' +                ", password='" + password + '\'' +                ", validationQuery='" + validationQuery + '\'' +                ", testOnBorrow=" + testOnBorrow +                '}';    }}

自此，多数据源动态切换的组件就搞完了，有木有赶脚身体被掏空。那么身体掏空被掏空后是谁在受益呢？当然是......

额，不要想多，我是说业务代码。来看下业务代码如何切换数据源：

  @Override    public HttpResult<Boolean> testMultiDataSource(UserCreateReqDTO userCreateReqDTO) {        if (userCreateReqDTO == null) {            return HttpResult.successResult(Boolean.FALSE);        }        UserDO userDO = UserConvent.conventToUserDO(userCreateReqDTO);        //先向默认数据源插入        if (!userDao.createUser(userDO)) {            throw new BusinessException(ErrorEnum.TEST_MULTI_DATASOURCE_EXCEPTION);        }        //再向起他数据源插入        List<DataSourceDO> dataSourceDOList = this.dataSourceDao.query();        for (DataSourceDO dataSourceDO : dataSourceDOList) {            DataSourceBeanBuilder builder = new DataSourceBeanBuilder(                    dataSourceDO.getDatabaseName(),                    dataSourceDO.getDatabaseIp(),                    dataSourceDO.getDatabasePort(),                    dataSourceDO.getDatasourceName(),                    dataSourceDO.getUsername(),                    dataSourceDO.getPassword());            DataSourceContext.setDataSource(builder);            if (!userDao.createUser(userDO)) {                throw new BusinessException(ErrorEnum.TEST_MULTI_DATASOURCE_EXCEPTION);            }            DataSourceContext.clearDataSource();        }        return HttpResult.successResult(Boolean.TRUE);    }

可以看到，当不适用DataSourceContext.setDataSource()方法设置数据源的时候，框架使用默认的数据源，即defaultDataSource参数配置的数据源。当时用DataSource.setDataSource()方法设置数据源之后，框架会使用指定的数据源。使用完毕后执行DataSource.clearDataSource()就又会切回到默认的数据源。

笔者已经有了一个是实现好的例子在我的github上，具体地址为点击打开链接。该工程是一个完整的ssm demo，并且其中包含了一些常用的组件，笔者还将继续增强他的功能。

觉得文章还不错的话，记得关注我哦！