一、基于Queue的数据源上下文实现
public class DataSourceContextHolder {// 使用ThreadLocal<Deque>支持嵌套数据源private static final ThreadLocal<Deque<DataSourceType>> CONTEXT_HOLDER = ThreadLocal.withInitial(LinkedList::new);/*** 设置数据源类型到队列头部*/public static void push(DataSourceType dataSourceType) {CONTEXT_HOLDER.get().push(dataSourceType);}/*** 获取当前数据源类型（队列头部）*/public static DataSourceType peek() {Deque<DataSourceType> deque = CONTEXT_HOLDER.get();return deque.isEmpty() ? null : deque.peek();}/*** 移除当前数据源类型*/public static DataSourceType pop() {Deque<DataSourceType> deque = CONTEXT_HOLDER.get();return deque.isEmpty() ? null : deque.pop();}/*** 清空数据源类型*/public static void clear() {CONTEXT_HOLDER.remove();}/*** 获取当前数据源切换深度*/public static int size() {return CONTEXT_HOLDER.get().size();}
}
二、数据源切换注解
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface DataSource {DataSourceType value() default DataSourceType.DB1;
}@Aspect
@Component
public class DataSourceAspect {@Pointcut("@annotation(com.example.annotation.DataSource)")public void dataSourcePointcut() {}@Around("dataSourcePointcut()")public Object around(ProceedingJoinPoint point) throws Throwable {MethodSignature signature = (MethodSignature) point.getSignature();Method method = signature.getMethod();DataSource ds = method.getAnnotation(DataSource.class);if (ds != null) {// 将数据源推入栈中DataSourceContextHolder.push(ds.value());}try {return point.proceed();} finally {if (ds != null) {// 方法执行完后弹出数据源DataSourceContextHolder.pop();}}}
}
三、使用示例
- 嵌套数据源服务
@Service
@Slf4j
public class NestedDataSourceService {@Autowiredprivate UserMapper userMapper;@Autowiredprivate OrderMapper orderMapper;@DataSource(DataSourceType.DB1)public User getUserFromDb1(Long id) {return userMapper.selectById(id);}@DataSource(DataSourceType.DB2)public Order getOrderFromDb2(Long id) {return orderMapper.selectById(id);}// 嵌套使用数据源的方法@DataSource(DataSourceType.DB1)public void complexOperation() {// 使用DB1数据源User user = getUserFromDb1(1L);log.info("Current DataSource: {}", DataSourceContextHolder.peek());// 切换到DB2数据源Order order = getOrderFromDb2(1L);log.info("Current DataSource: {}", DataSourceContextHolder.peek());// 自动恢复到DB1数据源user = getUserFromDb1(2L);log.info("Current DataSource: {}", DataSourceContextHolder.peek());}
}
- 事务处理
@Service
@Slf4j
public class NestedTransactionService {@Autowiredprivate DataSource db1DataSource;@Autowiredprivate DataSource db2DataSource;@Autowiredprivate PlatformTransactionManager db1TransactionManager;@Autowiredprivate PlatformTransactionManager db2TransactionManager;public void nestedTransactionOperation() {// DB1事务TransactionTemplate tt1 = new TransactionTemplate(db1TransactionManager);tt1.execute(status -> {DataSourceContextHolder.push(DataSourceType.DB1);try {// DB1操作// DB2事务TransactionTemplate tt2 = new TransactionTemplate(db2TransactionManager);tt2.execute(status2 -> {DataSourceContextHolder.push(DataSourceType.DB2);try {// DB2操作return null;} finally {DataSourceContextHolder.pop();}});return null;} finally {DataSourceContextHolder.pop();}});}
}
- 异常处理
@Service
@Slf4j
public class SafeDataSourceService {public <T> T executeWithDataSource(DataSourceType dataSourceType, Supplier<T> operation) {DataSourceContextHolder.push(dataSourceType);try {return operation.get();} finally {DataSourceContextHolder.pop();}}public void safeOperation() {try {// 使用DB1executeWithDataSource(DataSourceType.DB1, () -> {// DB1操作// 嵌套使用DB2return executeWithDataSource(DataSourceType.DB2, () -> {// DB2操作return null;});});} catch (Exception e) {log.error("操作失败", e);// 确保清理ThreadLocalDataSourceContextHolder.clear();}}
}
- 监控和调试
@Aspect
@Component
@Slf4j
public class DataSourceMonitorAspect {@Around("@annotation(dataSource)")public Object monitorDataSource(ProceedingJoinPoint point, DataSource dataSource) throws Throwable {String methodName = point.getSignature().getName();DataSourceType dsType = dataSource.value();log.debug("Method [{}] switching to datasource [{}], current stack depth: {}", methodName, dsType, DataSourceContextHolder.size());long startTime = System.currentTimeMillis();try {return point.proceed();} finally {long endTime = System.currentTimeMillis();log.debug("Method [{}] finished using datasource [{}], took {}ms", methodName, dsType, (endTime - startTime));}}
}
- 单元测试
@SpringBootTest
class NestedDataSourceTests {@Autowiredprivate NestedDataSourceService service;@Testvoid testNestedDataSourceSwitch() {// 验证数据源嵌套切换service.complexOperation();// 确保ThreadLocal被清理assertNull(DataSourceContextHolder.peek());assertEquals(0, DataSourceContextHolder.size());}@Testvoid testDataSourceStackOperation() {DataSourceContextHolder.push(DataSourceType.DB1);assertEquals(DataSourceType.DB1, DataSourceContextHolder.peek());DataSourceContextHolder.push(DataSourceType.DB2);assertEquals(DataSourceType.DB2, DataSourceContextHolder.peek());DataSourceContextHolder.pop();assertEquals(DataSourceType.DB1, DataSourceContextHolder.peek());DataSourceContextHolder.pop();assertNull(DataSourceContextHolder.peek());}
}
使用Queue（具体是Deque）来支持数据源嵌套使用的优点：
支持嵌套：可以处理方法嵌套调用时的数据源切换
自动恢复：使用栈结构自然地支持数据源的恢复
线程安全：ThreadLocal保证了线程隔离
易于调试：可以方便地查看数据源切换历史
异常安全：即使发生异常，也能保证数据源正确恢复
注意事项：
确保在finally块中调用pop()
考虑添加最大嵌套深度限制
在异常情况下清理ThreadLocal
监控数据源切换的性能影响
考虑添加数据源切换的日志记录
这种实现方式适合处理复杂的业务场景，特别是需要在同一个事务中访问多个数据源的情况。
