SpringBoot实现动态数据源切换的项目实践

<dependencies>
    <!-- Spring Boot Web 依赖 -->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-web</artifactId>
    </dependency>
    <!-- Spring Boot 数据源依赖 -->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-data-jpa</artifactId>
    </dependency>
    <!-- 数据库驱动依赖，以 MySQL 为例 -->
    <dependency>
        <groupId>mysql</groupId>
        <artifactId>mysql-connector-java</artifactId>
    </dependency>
</dependencies>

spring:
  datasource:
    primary:
      url: jdbc:mysql://localhost:3306/db1?useUnicode=true&characterEncoding=utf-8&serverTimezone=UTC
      username: root
      password: root
      driver-class-name: com.mysql.cj.jdbc.Driver
    secondary:
      url: jdbc:mysql://localhost:3306/db2?useUnicode=true&characterEncoding=utf-8&serverTimezone=UTC
      username: root
      password: root
      driver-class-name: com.mysql.cj.jdbc.Driver

@Configuration
public class DynamicDataSourceConfig {
    @Primary
    @Bean(name = "primaryDataSource")
    @ConfigurationProperties(prefix = "spring.datasource.primary")
    public DataSource dataSource1() {
        return DataSourceBuilder.create().build();
    }
    @Bean(name = "secondaryDataSource")
    @ConfigurationProperties(prefix = "spring.datasource.secondary")
    public DataSource dataSource2() {
        return DataSourceBuilder.create().build();
    }
}

public class DynamicDataSourceContextHolder {
    private static final ThreadLocal<String> CONTEXT_HOLDER = new ThreadLocal<>();
    public static void setDataSourceType(String dataSourceType) {
        CONTEXT_HOLDER.set(dataSourceType);
    }
    public static String getDataSourceType() {
        return CONTEXT_HOLDER.get();
    }
    public static void clearDataSourceType() {
        CONTEXT_HOLDER.remove();
    }
}

spring:
  datasource:
    type: com.zaxxer.hikari.HikariDataSource
    dynamic:
      primary: primaryDataSource
      secondary: secondaryDataSource

@Aspect
@Component
public class DynamicDataSourceAspect {
    @Before("@annotation(TargetDataSource)")
    public void switchDataSource(JoinPoint point) {
        // 获取方法上的注解
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        TargetDataSource ds = method.getAnnotation(TargetDataSource.class);
        // 判断注解是否为空
        if (ds == null) {
            return;
        }
        // 切换数据源
        DynamicDataSourceContextHolder.setDataSourceType(ds.name());
    }
    @After("@annotation(TargetDataSource)")
    public void restoreDataSource(JoinPoint point) {
        // 获取方法上的注解
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        TargetDataSource ds = method.getAnnotation(TargetDataSource.class);
        // 判断注解是否为空
        if (ds == null) {
            return;
        }
        // 恢复数据源
        DynamicDataSourceContextHolder.clearDataSourceType();
    }
}

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

@Service
public class UserService {
    @Autowired
    private UserMapper userMapper;
    @TargetDataSource(name = "primary")
    public List<User> findUsersFromPrimary() {
        return userMapper.findAll();
    }
    @TargetDataSource(name = "secondary")
    public List<User> findUsersFromSecondary() {
        return userMapper.findAll();
    }
}

@Entity
public class User {
    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;
    private String name;
    // 省略 getter 和 setter
}
public interface UserRepository extends JpaRepository<User, Long> {
}

@RunWith(SpringRunner.class)
@SpringBootTest
public class UserServiceTest {
    @Autowired
    private UserService userService;
    @Test
    public void testDynamicDataSource() {
        List<User> usersFromPrimary = userService.findUsersFromPrimary();
        List<User> usersFromSecondary = userService.findUsersFromSecondary();
        
        // 假设我们期望从主数据源和从数据源分别查询到不同的用户列表
        assertNotNull(usersFromPrimary);
        assertNotNull(usersFromSecondary);
        assertNotEquals(usersFromPrimary, usersFromSecondary);
    }
}