Spring Boot 动态数据源切换

背景

随着互联网应用的快速发展，多数据源的需求日益增多。Spring Boot 以其简洁的配置和强大的功能，成为实现动态数据源切换的理想选择。本文将通过具体的配置和代码示例，详细介绍如何在 Spring Boot 应用中实现动态数据源切换，帮助开发者高效应对不同业务场景下的数据管理需求。无论是读写分离还是数据隔离，都能轻松搞定。

AOP动态代理

AOP注解

@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
public @interface TargetDataSource {
    String name();
}

AOP切面类

@Aspect
@Component
public class DataSourceAspect {

    @Pointcut("@annotation(com.example.aliyunai.db.TargetDataSource)")
    public void dataSourcePointcut() {
    }

    @Before("dataSourcePointcut()")
    public void changeDataSource(JoinPoint point) {
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        TargetDataSource targetDataSource = method.getAnnotation(TargetDataSource.class);
        if (targetDataSource != null) {
            String dataSourceName = targetDataSource.name();
            DataSourceContextHolder.setDataSourceKey(dataSourceName);
        }
    }

    @After("dataSourcePointcut()")
    public void clearDataSource(JoinPoint point) {
        DataSourceContextHolder.clearDataSourceKey();
    }
}

数据源配置

@Configuration
public class DataSourceConfig {


    @Bean(name = "master")
    public DataSource primaryDataSource() {
        return DataSourceBuilder.create()
                .type(HikariDataSource.class)
                .driverClassName("")
                .url("")
                .username("")
                .password("")
                .build();
    }

    @Bean(name = "slave")
    public DataSource secondaryDataSource() {
        return DataSourceBuilder.create()
                .type(HikariDataSource.class)
                .driverClassName("")
                .url("")
                .username("")
                .password("")
                .build();
    }

    @Bean(name = "dynamicDataSource")
    @Primary
    public DynamicDataSource dynamicDataSource(@Qualifier("master") DataSource master,
                                               @Qualifier("slave") DataSource slave) {
        Map<Object, Object> targetDataSources = new HashMap<>();
        targetDataSources.put("master", master);
        targetDataSources.put("slave", slave);

        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        dynamicDataSource.setTargetDataSources(targetDataSources);
        dynamicDataSource.setDefaultTargetDataSource(master);
        return dynamicDataSource;
    }

}

线程上下文

public class DataSourceContextHolder {
    private static final ThreadLocal<String> contextHolder = new ThreadLocal<>();

    public static void setDataSourceKey(String key) {
        contextHolder.set(key);
    }

    public static String getDataSourceKey() {
        return contextHolder.get();
    }

    public static void clearDataSourceKey() {
        contextHolder.remove();
    }
}

动态数据源设置

public class DynamicDataSource extends AbstractRoutingDataSource {

    private static final Logger logger = LoggerFactory.getLogger(DynamicDataSource.class);

    @Override
    protected Object determineCurrentLookupKey() {
        String dataSourceKey = DataSourceContextHolder.getDataSourceKey();
        logger.info("Determining current data source: {}", dataSourceKey);
        return dataSourceKey;
    }
}

service类

@Service
public class UserService {

    @Resource
    private UserMapper userMapper;

    @TargetDataSource(name = "master")
    public User queryFromPrimary() {
        User user = userMapper.selectById(1);
        return user;
    }

    @TargetDataSource(name = "slave")
    public User queryFromSecondary() {
        User user = userMapper.selectById(1L);
        return user;
    }
}

Mybatis 拦截器

@Component
@Intercepts({
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class,Object.class}
        )})
public class DynamicDataSourceInterceptor implements Interceptor {

    private static final Logger logger = LoggerFactory.getLogger(DynamicDataSourceInterceptor.class);

    // 假设这里有一个获取当前数据源标识的方法，你需要根据实际项目中的实现来替换
    private String getCurrentDataSourceKey() {
        // 示例：从某个上下文持有者中获取数据源标识，这里只是示意，实际要替换
        return DataSourceContextHolder.getDataSourceKey();
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        MappedStatement mappedStatement = (MappedStatement) invocation.getArgs()[0];
        Object parameter = invocation.getArgs()[1];
        BoundSql boundSql = mappedStatement.getBoundSql(parameter);
        String sql = boundSql.getSql();
        //对单个表进行处理
        logger.info("sql:"+sql);
       
        String mappedStatementId = mappedStatement.getId();
        //对所有查询进行处理
        if (mappedStatementId.contains("query")) {
            String table = getTable(sql);
            if (table.equals("user")){
                DataSourceContextHolder.setDataSourceKey(getCurrentDataSourceKey());
            }
        //增删改
        }else if (mappedStatementId.contains("update")) {
            DataSourceContextHolder.setDataSourceKey(getCurrentDataSourceKey());
        }

        return invocation.proceed();
    }


    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }


    private static String getTable(String sql) {
        // 定义正则表达式模式，用于匹配 "from" 和 "where" 之间的表名
        Pattern pattern = Pattern.compile("FROM\\s+(\\w+)\\s+WHERE");
        Matcher matcher = pattern.matcher(sql);
        if (matcher.find()) {
            return matcher.group(1);
        }
        return null;
    }
}