Springboot+mybatis-plus+dynamic-datasource+继承DynamicRoutingDataSource切换数据源

背景

最近公司要求支持saas，实现动态切换库的操作，默认会加载主租户的数据源，其他租户数据源在使用过程中自动创建加入。

解决问题

1.通过请求中设置租户id 查询对应的库
2.通过设置上下文租户id 查询对应的库
3.测试mybatisplus mapper，service继承后设置上下文能否正常 查询对应的库

解决要求

1.改造现有系统尽量少改动,避免过多的耦合代码
2.已有功能正常
3.不影响之前的@DS注解切换数据源的

实现流程

1.代码结构

2.引入依赖

 <dependency>
            <groupId>com.alibaba</groupId>
            <artifactId>transmittable-thread-local</artifactId>
            <version>2.14.3</version>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-test</artifactId>
            <scope>test</scope>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>
        <dependency>
            <groupId>com.baomidou</groupId>
            <artifactId>mybatis-plus-boot-starter</artifactId>
            <version>3.5.3.1</version>
        </dependency>
        <dependency>
            <groupId>mysql</groupId>
            <artifactId>mysql-connector-java</artifactId>
            <version>5.1.47</version>
        </dependency>
        <dependency>
            <groupId>com.baomidou</groupId>
            <artifactId>dynamic-datasource-spring-boot-starter</artifactId>
            <version>4.2.0</version>
        </dependency>
        <dependency>
            <groupId>org.testng</groupId>
            <artifactId>testng</artifactId>
            <version>7.4.0</version>
        </dependency>
        <dependency>
            <groupId>org.projectlombok</groupId>
            <artifactId>lombok</artifactId>
        </dependency>

3.代码

3.1.TenantContextHolder

用于将租户id设置为上下文，获取当前的租户id

package com.liuhm.context;

import com.alibaba.ttl.TransmittableThreadLocal;

/**
 * saas 上下文 Holder
 */
public class TenantContextHolder {
    /**
     * 当前租户编号
     */
    private static final ThreadLocal<String> TENANT_ID = new TransmittableThreadLocal<>();


    /**
     * 获得租户编号。
     *
     * @return 租户编号
     */
    public static String getTenantId() {
        return TENANT_ID.get();
    }


    /**
     * 获得租户编号。如果不存在，则抛出 NullPointerException 异常
     *
     * @return 租户编号
     */
    public static String getRequiredTenantId() {
        String tenantId = getTenantId();
        if (tenantId == null) {
            throw new NullPointerException("TenantContextHolder 不存在租户编号！");
        }
        return tenantId;
    }

    public static void setTenantId(String tenantId) {
        TENANT_ID.set(tenantId);
    }

    public static void clear() {
        TENANT_ID.remove();
    }

}

3.2.TenantWebFilter

拦截所有的请求获取header或者url中租户id的值，然后设置到上下文中。

(获取租户id可以改成获取token，并将租户id存入token值中，方便获取租户id)

package com.liuhm.config;

import com.liuhm.context.TenantContextHolder;
import org.springframework.util.StringUtils;
import org.springframework.web.filter.OncePerRequestFilter;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.Objects;

public class TenantWebFilter extends OncePerRequestFilter {

    public static final String HEADER_TENANT_ID = "X-Tenant-Id";

    public static String getTenantId(HttpServletRequest request){
        String tenantId = StringUtils.hasLength(request.getHeader(HEADER_TENANT_ID)) ?
                request.getHeader(HEADER_TENANT_ID) :
                request.getHeader(HEADER_TENANT_ID.toLowerCase());
        if (StringUtils.isEmpty(tenantId)) {
            tenantId = getQueryParam(request.getQueryString(),HEADER_TENANT_ID);
        }
        return StringUtils.hasText(tenantId) ? tenantId : null;
    }

    public static String getQueryParam(String query,String key){
        if(Objects.isNull(query)){
            return null;
        }
        String[] params = query.split("&");
        for (String param : params) {
            String[] keyValue = param.split("=");
            if(Objects.equals(key.toLowerCase(),keyValue[0].toLowerCase()) && keyValue.length > 1){
                return keyValue[1];
            }
        }

        return null;
    }
    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain chain) throws ServletException, IOException{
        if (request.getRequestURI().equalsIgnoreCase("/harbor/clear")) {
            chain.doFilter(request, response);
        } else {
            String tenantId = getTenantId(request);
            if (tenantId != null) {
                TenantContextHolder.setTenantId(tenantId);
            }
            try {
                chain.doFilter(request, response);
            } finally {
                // 清理
                TenantContextHolder.clear();
            }
        }
    }
}

3.3 MyDynamicRoutingDataSource

• MyDynamicRoutingDataSource继承DynamicRoutingDataSource 重新修改选择数据源的逻辑。

• DynamicDataSourceContextHolder.peek()为空时，表示原功能默认的@DS没有设置，就通过tenantId去获取数据源

• getDataSourceProperty 通过tenantId 获取数据源的配置信息

• createDatasourceIfAbsent 通过配置信息去创建数据源并加入到dataSourceMap中

• 通过对应的key去获取对应的数据源

package com.liuhm.config;



import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.creator.DataSourceProperty;
import com.baomidou.dynamic.datasource.creator.DefaultDataSourceCreator;
import com.baomidou.dynamic.datasource.provider.DynamicDataSourceProvider;
import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.DynamicDataSourceProperties;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import com.liuhm.context.TenantContextHolder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Lazy;

import javax.annotation.Resource;
import javax.sql.DataSource;
import java.util.List;
import java.util.Set;


/**
 * @ClassName：MyDynamicRoutingDataSource
 * @Description: TODO
 * @Author: liuhaomin
 * @Date: 2024/5/9 8:44
 */
@Slf4j
public class MyDynamicRoutingDataSource extends DynamicRoutingDataSource {
    @Override
    public DataSource determineDataSource() {
        if(DynamicDataSourceContextHolder.peek() == null){
            String tenantId = TenantContextHolder.getTenantId();
            if(tenantId == null){
                throw new RuntimeException("租户id不能为空");
            }
            DataSourceProperty dataSourceProperty = getDataSourceProperty(tenantId);
            createDatasourceIfAbsent(dataSourceProperty);
            return getDataSource(tenantId);
        }else {
            DataSourceProperty dataSourceProperty = getDataSourceProperty(DynamicDataSourceContextHolder.peek());
            createDatasourceIfAbsent(dataSourceProperty);
            return super.determineDataSource();
        }
    }

    public MyDynamicRoutingDataSource(List<DynamicDataSourceProvider> providers) {
        super(providers);
    }



    /**
     * 用于创建租户数据源的 Creator
     */
    @Resource
    @Lazy
    private DefaultDataSourceCreator dataSourceCreator;

    @Resource
    @Lazy
    private DynamicDataSourceProperties dynamicDataSourceProperties;
    @Value("${spring.datasource.dynamic.primaryDatabase}")
    private String primaryDatabase;

    public DataSourceProperty getDataSourceProperty(String tenantId){

        DataSourceProperty dataSourceProperty = new DataSourceProperty();
        DataSourceProperty primaryDataSourceProperty = dynamicDataSourceProperties.getDatasource().get(dynamicDataSourceProperties.getPrimary());
        BeanUtils.copyProperties(primaryDataSourceProperty,dataSourceProperty);
        dataSourceProperty.setUrl(dataSourceProperty.getUrl().replace(primaryDatabase,tenantId));
        dataSourceProperty.setPoolName(tenantId);

        return dataSourceProperty;
    }

    private String createDatasourceIfAbsent(DataSourceProperty dataSourceProperty){
        // 1. 重点：如果数据源不存在，则进行创建
        if (isDataSourceNotExist(dataSourceProperty)) {
            // 问题一：为什么要加锁？因为，如果多个线程同时执行到这里，会导致多次创建数据源
            // 问题二：为什么要使用 poolName 加锁？保证多个不同的 poolName 可以并发创建数据源
            // 问题三：为什么要使用 intern 方法？因为，intern 方法，会返回一个字符串的常量池中的引用
            // intern 的说明，可见 https://www.cnblogs.com/xrq730/p/6662232.html 文章
            synchronized(dataSourceProperty.getPoolName().intern()){
                if (isDataSourceNotExist(dataSourceProperty)) {
                    log.debug("创建数据源：{}", dataSourceProperty.getPoolName());
                    DataSource dataSource = null;
                    try {
                        dataSource = dataSourceCreator.createDataSource(dataSourceProperty);
                    }catch (Exception e){
                        log.error("e {}",e);
                        if(e.getMessage().contains("Unknown database")){
                            throw new RuntimeException("租户不存在");
                        }
                        throw e;
                    }
                    addDataSource(dataSourceProperty.getPoolName(), dataSource);
                }
            }
        } else {
            log.debug("数据源已存在,无需创建：{}", dataSourceProperty.getPoolName());
        }
        // 2. 返回数据源的名字
        return dataSourceProperty.getPoolName();
    }

    private boolean isDataSourceNotExist(DataSourceProperty dataSourceProperty){
        return !getDataSources().containsKey(dataSourceProperty.getPoolName());
    }
}

3.4.TenantAutoConfiguration

• TenantWebFilter加入FilterRegistrationBean
• 创建 MyDynamicRoutingDataSource Bean

package com.liuhm.config;

import com.baomidou.dynamic.datasource.provider.DynamicDataSourceProvider;
import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.DynamicDataSourceProperties;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.web.servlet.FilterRegistrationBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.sql.DataSource;
import java.util.List;

@Configuration
public class TenantAutoConfiguration {

    @Bean
    public FilterRegistrationBean<TenantWebFilter> tenantContextWebFilter() {
        FilterRegistrationBean<TenantWebFilter> registrationBean = new FilterRegistrationBean<>();
        registrationBean.setFilter(new TenantWebFilter());
        registrationBean.setOrder(-104);
        return registrationBean;
    }

    @Autowired
    private DynamicDataSourceProperties properties;
    @Bean
    public DataSource dataSource(List<DynamicDataSourceProvider> providers) {
        MyDynamicRoutingDataSource dataSource = new MyDynamicRoutingDataSource(providers);
        dataSource.setPrimary(properties.getPrimary());
        dataSource.setStrict(properties.getStrict());
        dataSource.setStrategy(properties.getStrategy());
        dataSource.setP6spy(properties.getP6spy());
        dataSource.setSeata(properties.getSeata());
        dataSource.setGraceDestroy(properties.getGraceDestroy());
        return dataSource;
    }
}

4.总结

4.1.多租户切换的方法

1. dynamic-datasource 跨库进行切换数据源可以用DynamicDataSourceContextHolder.push()

• 在过滤器[filter]里切换
• 拦截器里切换数据源
• 方法内部硬编码切换
• 通过service，mapper加注解进行切换@DS （不推荐，有切面没有切成功的，如本类调用自己的方法）

2. 重写DynamicRoutingDataSource选择器，自定义上下文获取租户id获取对应的DataSource

4.2.上诉方法中都可以实现

• 过滤器和拦截器切换数据源的时候，线程执行的方法不容切换，需要手动切换，或者在设置租户id的时候进行切换数据源。（耦合性过大，代码不够单一，如果在设置租户id的时候去切换数据源）
• 重写DynamicRoutingDataSource选择器，只是在执行sql前进行数据源获取的切换，耦合性小，代码单一性好，且不影响之前的功能。

4.3.设置租户id需要注意的

• 所有请求需要拦截进行设置
• 所有线程需要相关的需要进行重写并设置租户上下文
• 所有fegin需要进行设置租户上下文
• 以上4.3的操作可以学习一下mdc链路追踪日志的代码

编码不易，有问题多多指教

博客地址

代码下载

下面的springcloud_dynamic_datasource_tenant