SpringBoot 实现动态数据源

功能：前端请求接口时携带用户信息，后端拦截器获取用户信息后切换数据源查询数据。

使用场景：多租户，sass，pass等项目。

实现原理：主要通过SpringBoot提供的AbstractRoutingDataSource类，该类允许我们存入一个map，然后在通过数据源获取数据前允许我们指定一个key，然后自动的根据key从map内取对应的数据源。

1 创建 DataSourceHolder类管理数据源的创建和存储

@Component
@NoArgsConstructor
@Slf4j
public class DataSourceHolder {

    @Autowired
    private AbstractRoutingDataSource routingDataSource;

    @Resource
    private DataSourceInfoMapper dataSourceInfoMapper;

    /**
     * 存放所有数据源,其实 key = dataSourceKey , value = DataSource
     */
    protected final ConcurrentHashMap<Object, Object> dataSourceMap = new ConcurrentHashMap<>(64);

    /**
     * 用于初始化数据源时加锁
     */
    private final ConcurrentHashMap<String, Object> dataSourceInitLock = new ConcurrentHashMap<>(64);

    /**
     * 保存当前的 数据源key
     */
    private static ThreadLocal<String> currDataSourceKey = new ThreadLocal<>();

    /**
     * 是否使用默认的数据源
     */
    private static ThreadLocal<Boolean> useDefaultDs = new ThreadLocal<>();

    public static void setDataSourceKey(String dataSourceKey){
        log.info("数据源切换:{} -> {}", currDataSourceKey.get(), dataSourceKey  );
        currDataSourceKey.set(dataSourceKey);
        unUseDefaultDataSource();
    }

    public static void useDefaultDataSource(){
        useDefaultDs.set(true);
    }
    public static void unUseDefaultDataSource(){
        useDefaultDs.remove();
    }

    /**
     * 获取当前的数据源 key
     * @return 当前 dataSourcekey
     */
    public String getDataSourceKey(){

        //获取指定的数据源 KEY
        final String currentDsKey = currDataSourceKey.get();

        //如果指定使用默认数据源 或者 未指定数据源，返回null
        if( Objects.equals(useDefaultDs.get(),true) || StrUtil.isBlank( currentDsKey ) ){
            log.info("采用默认数据源");
            return null;
        }

        //如果指定的数据源已经被初始化
        if( dataSourceMap.containsKey(currentDsKey) ){
            log.info("采用动态数据源,dataSourceKey:{}",currentDsKey);
            return currentDsKey;
        }

        //如果数据源未被初始化，尝试初始化数据源
        return initDataSource( currentDsKey );
    }


    /**
     * 初始化数据源
     * @return 返回数据源key
     */
    private String initDataSource( final String dataSourceKey ) {

        //如果不存在，设置一个对象，主要用于加锁
        Object lockObj = dataSourceInitLock.computeIfAbsent(dataSourceKey, key -> new Object());


        synchronized ( lockObj ){
            log.info("开始初始化数据源,dataSourceKey:{}",dataSourceKey);

            if( dataSourceMap.containsKey(dataSourceKey) ){
                log.info("数据源已被其他线程初始化,dataSourceKey:{}",dataSourceKey);
                return dataSourceKey;
            }

            try {
                //使用默认的数据源
                useDefaultDataSource();

                //根据dataSourceKey,查询数据库配置信息
                LambdaQueryWrapper<DataSourceInfo> wrapper = Wrappers.lambdaQuery();
                wrapper.eq(DataSourceInfo::getDataSourceKey,dataSourceKey);
                wrapper.last( "limit 1");

                DataSourceInfo dataSourceInfo = dataSourceInfoMapper.selectOne(wrapper);

                if( dataSourceInfo == null ){
                    //这里要抛异常，否则会使用默认的数据源
                    log.error("加载数据源失败，数据源配置信息不存在，dataSourceKey:{}",dataSourceKey);
                    return null;
                }

                Class dataSourceType = Class.forName( dataSourceInfo.getDataSourceType() );

                //创建数据源
                DataSource dataSource = DataSourceBuilder.create()
                    .driverClassName(dataSourceInfo.getDriverClassName())
                    .username(dataSourceInfo.getUsername())
                    .password(dataSourceInfo.getPassword())
                    .url(dataSourceInfo.getUrl())
                    .type(dataSourceType)
                    .build();

                //放入数据源map
                dataSourceMap.put(dataSourceKey,dataSource);
                //数据源改变后需要刷新
                refreshDataSource();

                //创建成功后返回
                return dataSourceKey;
            }catch( ClassNotFoundException ignored){
            } finally {
                unUseDefaultDataSource();
            }

        }
        return null;
    }

    /**
     * 删除数据源
     * @param dataSourceKey 数据源Key
     */
    public void deleteDataSource( String dataSourceKey ){
        dataSourceMap.remove(dataSourceKey);
        refreshDataSource();
    }


    /**
     * 刷新数据源
     */
    public void refreshDataSource(){
        log.info("刷新动态数据源");
        routingDataSource.afterPropertiesSet();
    }

}

2 继承AbstractRoutingDataSource实现动态数据源

@NoArgsConstructor
public class DynamicDataSource extends AbstractRoutingDataSource {


    @Resource
    private DataSourceHolder dataSourceHolder;


    /**
     * 主要用于获取当前使用的数据源的key
     * @return
     */
    @Override
    protected Object determineCurrentLookupKey() {
        return dataSourceHolder.getDataSourceKey();
    }

}

3 配置数据源

@Configuration
@AutoConfigureAfter(DataSourceAutoConfiguration.class)
public class DynamicDataSourceConfiguration {

    @Autowired
    private DataSourceHolder dataSourceHolder;


    @Bean("defaultDataSource")
    public DataSource defaultDataSource(DataSourceProperties dataSourceProperties){

        //创建数据源
        return DataSourceBuilder.create()
                .driverClassName(dataSourceProperties.getDriverClassName())
                .username(dataSourceProperties.getUsername())
                .password(dataSourceProperties.getPassword())
                .url(dataSourceProperties.getUrl())
                .type(dataSourceProperties.getType())
                .build();
    }

    @Primary
    @Bean("dataSource")
    public DynamicDataSource dynamicDataSource(){
        DynamicDataSource dynamicDataSource = new DynamicDataSource();

        // 设置默认的数据源
        dynamicDataSource.setDefaultTargetDataSource( SpringUtil.getBean("defaultDataSource") );
        //设置数据源map
        dynamicDataSource.setTargetDataSources( dataSourceHolder.dataSourceMap );

        return dynamicDataSource;
    }

    /**
     * 配置@Transactional注解事物
     */
    @Bean
    public PlatformTransactionManager transactionManager() {
        return new DataSourceTransactionManager(dynamicDataSource());
    }

}

4 创建过滤器根据请求动态切换数据源

@Configuration
public class WebConfiguration implements WebMvcConfigurer {


    /**
     * 注册filter
     */
    @Bean
    public FilterRegistrationBean filterRegistrationBean(){

        FilterRegistrationBean bean = new FilterRegistrationBean();

        bean.setFilter((servletRequest, servletResponse, filterChain) -> {
            //默认使用默认的数据源
            DataSourceHolder.useDefaultDataSource();

            Object dsKey = servletRequest.getParameter("dsKey");
            if( dsKey!=null ){
                //如果指定了数据源，使用指定的数据源
                DataSourceHolder.setDataSourceKey(dsKey.toString());
            }
            filterChain.doFilter(servletRequest,servletResponse);
        });
        bean.setOrder(1);
        bean.addUrlPatterns("/*");
        return bean;
    }

}