基于Spring JDBC AbstractRoutingDataSource 实现动态数据源

AbstractRoutingDataSource 实现动态数据源

AbstractRoutingDataSource 即抽象的路由数据源，提供了动态数据源切换的机制。你可以通过实现它的 determineCurrentLookupKey() 方法，根据不同的条件返回对应的数据源 key，基于这点可以根据外部输入完成数据源的动态选择

路由数据源实现

基于ThreadLocal实现，每次选择数据源都根据ThreadLocal中的压测标识来决定是否走正常数据源还是压测数据源

package com.huakai.springenv.perf;import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;public class TrafficRoutingDataSource extends AbstractRoutingDataSource {// ThreadLocal 用于存储当前请求是否是压测流量private static final ThreadLocal<Boolean> PRESSURE_TEST_FLAG = new ThreadLocal<>();// 设置压测标识public static void setPressureTestFlag(boolean isPressureTest) {PRESSURE_TEST_FLAG.set(isPressureTest);}// 清除压测标识public static void clearPressureTestFlag() {PRESSURE_TEST_FLAG.remove();}@Overrideprotected Object determineCurrentLookupKey() {// 根据 ThreadLocal 判断是否是压测流量Boolean isPressureTest = PRESSURE_TEST_FLAG.get();if (Boolean.TRUE.equals(isPressureTest)) {return "shadowDataSource";  // 返回影子库的数据源 key} else {return "primaryDataSource";  // 返回原始库的数据源 key}}
}

数据源配置

这里配置TrafficRoutingDataSource数据源，使用@Primary声明为Spring的默认数据源，并且将primaryDataSource和shadowDataSource添加到路由数据源列表

package com.huakai.springenv.perf;import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;@Configuration
public class DataSourceConfig {@Bean@Primarypublic DataSource routingDataSource(@Qualifier("primaryDataSource") DataSource primaryDataSource,@Qualifier("shadowDataSource") DataSource shadowDataSource) {TrafficRoutingDataSource routingDataSource = new TrafficRoutingDataSource();// 配置数据源Map<Object, Object> dataSourceMap = new HashMap<>();dataSourceMap.put("primaryDataSource", primaryDataSource);  // 主库dataSourceMap.put("shadowDataSource", shadowDataSource);    // 影子库routingDataSource.setTargetDataSources(dataSourceMap);routingDataSource.setDefaultTargetDataSource(primaryDataSource);  // 默认走主库routingDataSource.afterPropertiesSet();return routingDataSource;}// 配置主数据源和影子数据源@Beanpublic DataSource primaryDataSource() {HikariConfig hikariConfig = new HikariConfig();hikariConfig.setDriverClassName("com.mysql.cj.jdbc.Driver");hikariConfig.setJdbcUrl("jdbc:mysql://localhost:3306/test");  // 主库的 JDBC URLhikariConfig.setUsername("root");  // 数据库用户名hikariConfig.setPassword("123456");  // 数据库密码hikariConfig.setMaximumPoolSize(10);  // 最大连接池数hikariConfig.setMinimumIdle(0);return new HikariDataSource(hikariConfig);}@Beanpublic DataSource shadowDataSource() {HikariConfig hikariConfig = new HikariConfig();hikariConfig.setDriverClassName("com.mysql.cj.jdbc.Driver");hikariConfig.setJdbcUrl("jdbc:mysql://localhost:3306/perf__test");  // 从库的 JDBC URLhikariConfig.setUsername("root");  // 数据库用户名hikariConfig.setPassword("123456");  // 数据库密码hikariConfig.setMaximumPoolSize(10);  // 最大连接池数hikariConfig.setMinimumIdle(0);return new HikariDataSource(hikariConfig);}
}

基于过滤器实现压测流量自动选择影子库

过滤器

package com.huakai.springenv.perf;import org.springframework.stereotype.Component;import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.FilterConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import java.io.IOException;
@Component
public class PressureTestFilter implements Filter {@Overridepublic void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)throws IOException, ServletException {try {HttpServletRequest request1 = (HttpServletRequest) request;// 通过请求中的 header 判断是否为压测流量String pressureTestHeader = request1.getHeader("X-Pressure-Test");if ("true".equals(pressureTestHeader)) {TrafficRoutingDataSource.setPressureTestFlag(true);}chain.doFilter(request, response);} finally {// 清理 ThreadLocalTrafficRoutingDataSource.clearPressureTestFlag();}}@Overridepublic void init(FilterConfig filterConfig) throws ServletException {}@Overridepublic void destroy() {}
}

过滤器配置

package com.huakai.springenv.perf;import org.springframework.boot.web.servlet.FilterRegistrationBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;@Configuration
public class WebConfig {@Beanpublic FilterRegistrationBean testFilterRegistration() {FilterRegistrationBean registration = new FilterRegistrationBean(new PressureTestFilter());registration.addUrlPatterns("/"); //registration.setName("perfFilter");return registration;}
}

web入口

package com.huakai.springenv.perf;import com.huakai.springenv.mapper.KvMapper;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;import javax.annotation.Resource;@RestController()
@RequestMapping("perf")
public class PerfController {@Resourceprivate KvMapper kvMapper;@RequestMapping(value = "/test")public String test(@RequestParam String key) {return kvMapper.get(key).getValue();}
}