Spring Boot 实现多数据源切换

前言

Spring 本身支持多种数据源的核心类是 org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource, 这个类的核心方法有 3 个, 每个方法的作用如下

setTargetDataSources(Map<Object, Object> targetDataSources) 设置多数据源
setDefaultTargetDataSource(Object defaultTargetDataSource) 设置默认数据源, 当无法确认数据源时, 会启用这个默认数据源
determineCurrentLookupKey() 决定使用哪个数据源

基于这个类,实现 Spring 动态切换数据源的大致思路如下:

首先配置多个数据源
实现 AbstractRoutingDataSource 的三个核心方法, 可以根据我们的逻辑, 动态设置数据源
自定义注解, 设置每个方法需要的数据源
使用 AOP 实现数据源的动态切换

AbstractRoutingDataSource 实现逻辑

package com.baidu.sws.multidatasource;

import java.util.Stack;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class DbSourceContext {

    /*
     * 通过 ThreadLocal 传递当前所需的数据源
     */
    private static final ThreadLocal<String> TYPE_DB_CONTEXT = new ThreadLocal<String>();
    private static final ThreadLocal<Stack<String>> METHOD_DB_CONTEXT = new ThreadLocal<>();

    /**
     * 设置类级别数据源
     * 
     * @param source 数据源
     */
    public static void setTypeDbSource(String source) {
        log.debug("set source ====>" + source);
        TYPE_DB_CONTEXT.set(source);
    }

    /**
     * 设置方法级别数据源
     * 
     * @param source 数据源
     */
    public static void setMethodDbSource(String source) {
        log.debug("set method source ====>" + source);
        if (METHOD_DB_CONTEXT.get() == null) {
            METHOD_DB_CONTEXT.set(new Stack<>());
        }
        METHOD_DB_CONTEXT.get().push(source);
    }

    /**
     * 获取当前所需的数据源
     * 方法级别 > 类级别
     */
    public static String getDbSource() {

        String dbName = METHOD_DB_CONTEXT.get() == null ? TYPE_DB_CONTEXT.get() : getMethodDbSource();
        log.debug("get source ====>" + dbName);
        return dbName;
    }

    /**
     * 获取类级别数据源
     * 
     * @return
     */
    public static String getMethodDbSource() {
        if (METHOD_DB_CONTEXT.get() == null || METHOD_DB_CONTEXT.get().isEmpty()) {
            log.debug("get method source ====> null" );
            return null;
        }
        return METHOD_DB_CONTEXT.get().peek();
    }

    /**
     * 清理类级别数据源
     * 
     */
    public static void clearTypeDbSource() {
        TYPE_DB_CONTEXT.remove();
    }

    /**
     * 清理方法级别数据源
     */
    public static void clearMethodDbSource() {
        if (METHOD_DB_CONTEXT.get() != null) {
            METHOD_DB_CONTEXT.get().pop();
            if (METHOD_DB_CONTEXT.get().isEmpty()) {
                METHOD_DB_CONTEXT.remove();
            }
        }
    }

}

注意：这里使用栈的数据结构存储数据源名称，这样可以适配在不同方法中使用不同的数据源。

例如如下场景
public class TestDBSourceService {

    @DBSource("db1")
    public void testDB1(){
        db1Dao.selectById(1);
        testDB2();
        db1Dao.selectById(2);
    }

    @DBSource("db2")
    public void testDB2(){
        db2Dao.selectById(1);
    }

}
首先在
执行 testDB1方法是栈内会被压入 `db1` 数据源，
当调用 `testDB2` 时栈被压入 `db2` 数据源，执行完后 `db2` 被弹出栈，此时栈顶为 `db1`
当继续执行 `db1Dao.selectById(2)` 时, 使用 `db1` 数据源

public class DynamicDataSource extends AbstractRoutingDataSource {

    /**
    * 从 DbSourceContext 中获取所需的数据源
    */
    @Override
    protected Object determineCurrentLookupKey() {
        return DbSourceContext.getDbSource();
    }
}

AOP 动态切换数据源实现逻辑

自定义注解

@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface DbSource {
    
    /**
     * 指定 Mapper 使用的数据源, 支持类或者特定的方法
     * @return -
     */
    String value();

}

AOP 动态切换数据源

import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.After;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import com.baidu.sws.annotation.DbSource;

@Aspect
@Order(-1) // 该切面应当先于 @Transactional 执行
@Component
public class DynamicDataSourceAspect {

    /**
     * 拦截有@DbSource注解的方法
     * @param joinPoint -
     * @return -
     * @throws Throwable -
     */
    @Around("@annotation(com.baidu.sws.annotation.DbSource)")
    public Object intercept(ProceedingJoinPoint joinPoint) throws Throwable {
        // 获取拦截的方法
        Signature signature = joinPoint.getSignature();
        if (signature == null || !(signature instanceof MethodSignature)) {
            return joinPoint.proceed();
        }
        // 获取注解
        MethodSignature methodSignature = (MethodSignature) signature;
        DbSource dbSource = methodSignature.getMethod().getAnnotation(DbSource.class);
        if (dbSource == null) {
            return joinPoint.proceed();
        }
        // 切换数据源
        DbSourceContext.setMethodDbSource(dbSource.value());
        Object result = joinPoint.proceed();
        // 将数据源置为默认数据源
        DbSourceContext.clearMethodDbSource();
        return result;
    }

    /**
     * 切换数据源
     * 
     * @param point    -
     * @param dbSource -
     */
    // 注解在类对象，拦截有@DbSource类下所有的方法
    @Before("@within(dbSource)")
    public void switchDataSource(JoinPoint point, DbSource dbSource) {
        // 切换数据源
        DbSourceContext.setTypeDbSource(dbSource.value());
    }

    /**
     * 重置数据源
     * 
     * @param point    -
     * @param dbSource -
     */
    // 注解在类对象，拦截有@DbSource类下所有的方法
    @After("@within(dbSource)")
    public void restoreDataSource(JoinPoint point, DbSource dbSource) {
        // 将数据源置为默认数据源
        DbSourceContext.clearTypeDbSource();
    }
}

多数据源配置

Spring 的 application.yaml 配置

spring:
  datasource:
    admin:
      driver-class-name: com.mysql.cj.jdbc.Driver
      password: *****
      jdbc-url: *****
      username: root

    books:
      jdbc-url: jdbc:sqlite:C:/Users/PC/Desktop/books.db
      driver-class-name: org.sqlite.JDBC

多数据源配置文件

@Configuration
public class DynamicDatasource {

    /**
     * mybatis type alias
     */
    private final String typeAliasesPackage = "cn.yogehaoren.ironblog_manage.entity";

    @Bean(name = "adminDataSource")
    @ConfigurationProperties(prefix = "spring.datasource.admin")
     public DataSource adminDataSource() {
        return DataSourceBuilder.create().build();
    }

    @Bean(name = "booksDataSource")
    @ConfigurationProperties(prefix = "spring.datasource.books")
    public DataSource booksDataSource() {
        return DataSourceBuilder.create().build();
    }


    @Bean(name = "haloDataSource")
    @ConfigurationProperties(prefix = "spring.datasource.halo")
    public DataSource haloDataSource() {
        return DataSourceBuilder.create().build();
    }


    @Bean(name = "dynamicDataSource")
    public DataSource dynamicDataSource() {
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        //配置多数据源
        Map<Object, Object> dbMap = new HashMap<>();
        dbMap.put("booksDataSource", booksDataSource());
        dbMap.put("haloDataSource", haloDataSource());
        dbMap.put("adminDataSource", adminDataSource());
        dynamicDataSource.setTargetDataSources(dbMap);

        // 设置默认数据源
        dynamicDataSource.setDefaultTargetDataSource(adminDataSource());

        return dynamicDataSource;
    }

    /*
     * 数据库连接会话工厂
     * 将动态数据源赋给工厂
     * mapper存于resources/mapper目录下
     * 默认bean存于com.main.example.bean包或子包下，也可直接在mapper中指定
     */
    @Bean(name = "sqlSessionFactory")
    public SqlSessionFactoryBean sqlSessionFactory() throws Exception {
        SqlSessionFactoryBean sqlSessionFactory = new SqlSessionFactoryBean();
        sqlSessionFactory.setDataSource(dynamicDataSource());
        sqlSessionFactory.setConfigLocation(new ClassPathResource("mybatis-configuration.xml"));
        sqlSessionFactory.setTypeAliasesPackage(typeAliasesPackage); //扫描bean
        PathMatchingResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
        sqlSessionFactory.setMapperLocations(resolver.getResources("classpath*:mapper/*.xml"));    // 扫描映射文件
        return sqlSessionFactory;
    }

    /**
    @Primary 如果使用 Mybatis Plus 需要使用这个Sql Bean
    @Bean(name = "sqlSessionFactory")
    public SqlSessionFactory sqlSessionFactory() throws Exception {
        MybatisSqlSessionFactoryBean sqlSessionFactory = new MybatisSqlSessionFactoryBean();
        sqlSessionFactory.setDataSource(dynamicDatasource());
        sqlSessionFactory.setTypeAliasesPackage(MODEL_PACKAGE); // 扫描bean
        PathMatchingResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
        sqlSessionFactory.setMapperLocations(resolver.getResources("classpath*:mapper/*.xml")); // 扫描映射文件
        return sqlSessionFactory.getObject();
    }
    */
    @Bean
    public PlatformTransactionManager transactionManager() {
        // 配置事务管理, 使用事务时在方法头部添加@Transactional注解即可
        return new DataSourceTransactionManager(dynamicDataSource());
    }

}

主函数

@SpringBootApplication(exclude = {DataSourceAutoConfiguration.class}) // 去除默认数据源配置
@MapperScan("com.example.dao")  // dao 类全类名
@EnableAspectJAutoProxy // 开启 aop
public class Main {
    
    public static void main(String[] args) {
        ConfigurableApplicationContext context = SpringApplication.run(Main.class, args);
    }

}

使用示例

Mapper 文件

@Mapper
@DbSource("booksDataSource")
public interface BookAdminMapper {

    @Select("Select * from admin_user where id = 1")
    BookAdmin getAdmin();

}

Spring Boot 实现多数据源切换

前言

AbstractRoutingDataSource 实现逻辑

AOP 动态切换数据源实现逻辑

多数据源配置

主函数

使用示例

评论

Golang 上传文件至阿里云盘

使用Flink DataStream 导入 Doris 数据

Spring boot logback配置文件保留spring boot默认日志样式

使用 GO 签发 Let's encrypt 免费证书

在编译器中调试mitmproxy脚本

Quartz 每个任务日志单独保存至数据库中

Retrofit2 集成 Spring

Spring Boot 实现多数据源切换

深度解析 Mybatis-Spring Mapper文件如何支持自动注入

Nginx 根据域名分发不同证书

狐狸的搬砖生活