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;}@Overrideprotected 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 {@Overridepublic 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@Lazyprivate DefaultDataSourceCreator dataSourceCreator;@Resource@Lazyprivate 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 {@Beanpublic FilterRegistrationBean<TenantWebFilter> tenantContextWebFilter() {FilterRegistrationBean<TenantWebFilter> registrationBean = new FilterRegistrationBean<>();registrationBean.setFilter(new TenantWebFilter());registrationBean.setOrder(-104);return registrationBean;}@Autowiredprivate DynamicDataSourceProperties properties;@Beanpublic 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.多租戶切換的方法

dynamic-datasource 跨庫進行切換數據源可以用DynamicDataSourceContextHolder.push()

在過濾器[filter]里切換
攔截器里切換數據源
方法內部硬編碼切換
通過service，mapper加注解進行切換@DS （不推薦，有切面沒有切成功的，如本類調用自己的方法）

重寫DynamicRoutingDataSource選擇器，自定義上下文獲取租戶id獲取對應的DataSource

4.2.上訴方法中都可以實現

過濾器和攔截器切換數據源的時候，線程執行的方法不容切換，需要手動切換，或者在設置租戶id的時候進行切換數據源。（耦合性過大，代碼不夠單一，如果在設置租戶id的時候去切換數據源）
重寫DynamicRoutingDataSource選擇器，只是在執行sql前進行數據源獲取的切換，耦合性小，代碼單一性好，且不影響之前的功能。