SpringBatch中如何在JPARepository内切换多数据源？

基于你已实现的AbstractRoutingDataSource和数据源上下文持有类，以下是具体实现方案，核心是通过ThreadLocal上下文切换数据源，并确保线程安全和事务兼容性：

1. 完善数据源上下文持有类 #

确保你的DataSourceContextHolder具备保存当前数据源、切换、恢复的能力，避免线程池复用导致的数据源污染：

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

    // 设置当前数据源key（比如"db2"、"postgres"）
    public static void setDataSourceKey(String key) {
        CONTEXT.set(key);
    }

    // 获取当前数据源key
    public static String getDataSourceKey() {
        return CONTEXT.get();
    }

    // 清除上下文
    public static void clearDataSourceKey() {
        CONTEXT.remove();
    }
}

2. 在JPA Repository自定义方法中切换数据源 #

由于JPA Repository的自定义方法需要实现类，你可以直接在实现逻辑中手动切换数据源，或者用AOP简化操作：

方式一：手动切换（直观可控） #

先定义Repository接口：

public interface ClientRepository extends JpaRepository<Client, Long> {
    // 从DB2读取客户端数据的方法
    List<Client> findByClientCodeIn(List<String> clientCodes);

    // 自定义方法：从Postgres读取好友数据
    List<Friend> findAllWithFriends(List<String> clientCodes);
}

编写自定义方法的实现类（命名规则为接口名+Impl，Spring Data JPA会自动识别）：

@Repository
public class ClientRepositoryImpl implements ClientRepository {
    @PersistenceContext
    private EntityManager entityManager;

    @Override
    public List<Friend> findAllWithFriends(List<String> clientCodes) {
        // 保存当前数据源，后续恢复
        String originalKey = DataSourceContextHolder.getDataSourceKey();
        try {
            // 切换到Postgres数据源
            DataSourceContextHolder.setDataSourceKey("postgres");
            // 执行JPQL查询（或原生SQL）
            return entityManager.createQuery(
                    "SELECT f FROM Friend f WHERE f.clientCode IN :codes", Friend.class)
                    .setParameter("codes", clientCodes)
                    .getResultList();
        } finally {
            // 恢复原数据源，避免线程污染
            if (originalKey != null) {
                DataSourceContextHolder.setDataSourceKey(originalKey);
            } else {
                DataSourceContextHolder.clearDataSourceKey();
            }
        }
    }

    // 继承的其他方法（如findByClientCodeIn）会使用默认数据源DB2
    @Override
    public List<Client> findByClientCodeIn(List<String> clientCodes) {
        return entityManager.createQuery(
                "SELECT c FROM Client c WHERE c.clientCode IN :codes", Client.class)
                .setParameter("codes", clientCodes)
                .getResultList();
    }
}

方式二：AOP切面切换（解耦） #

如果不想在Repository实现中写切换逻辑，可以用AOP拦截自定义方法自动切换：

@Aspect
@Component
public class DataSourceSwitchAspect {
    // 拦截自定义的findAllWithFriends方法
    @Around("execution(* com.yourpackage.ClientRepository.findAllWithFriends(..))")
    public Object switchToPostgres(ProceedingJoinPoint point) throws Throwable {
        String originalKey = DataSourceContextHolder.getDataSourceKey();
        try {
            DataSourceContextHolder.setDataSourceKey("postgres");
            return point.proceed();
        } finally {
            // 恢复原数据源
            if (originalKey != null) {
                DataSourceContextHolder.setDataSourceKey(originalKey);
            } else {
                DataSourceContextHolder.clearDataSourceKey();
            }
        }
    }
}

3. 自定义RepositoryItemReader中整合双数据源数据 #

在你的自定义RepositoryItemReader中，先读取DB2的客户端数据，再调用自定义方法拉取Postgres的好友数据，一次性组装成目标对象：

public class ClientWithFriendsItemReader extends RepositoryItemReader<ClientWithFriends> {
    private final ClientRepository clientRepository;

    public ClientWithFriendsItemReader(ClientRepository clientRepository) {
        this.clientRepository = clientRepository;
        // 配置Reader基础参数：读取DB2客户端数据的方法和参数
        setRepository(clientRepository);
        setMethodName("findByClientCodeIn");
        setArguments(Collections.singletonList(getBatchClientCodes())); // 传入5000个客户端编码
        // 可选：配置分页、排序
        setPageSize(5000);
    }

    @Override
    protected List<ClientWithFriends> doPageRead() {
        // 1. 从DB2读取客户端数据
        List<Client> clients = (List<Client>) super.doPageRead();
        if (clients.isEmpty()) {
            return Collections.emptyList();
        }

        // 2. 提取客户端编码，调用自定义方法从Postgres拉取好友数据
        List<String> clientCodes = clients.stream()
                .map(Client::getClientCode)
                .collect(Collectors.toList());
        List<Friend> friends = clientRepository.findAllWithFriends(clientCodes);

        // 3. 组装客户端与好友数据
        Map<String, List<Friend>> friendMap = friends.stream()
                .collect(Collectors.groupingBy(Friend::getClientCode));
        return clients.stream().map(client -> {
            ClientWithFriends result = new ClientWithFriends();
            result.setClient(client);
            result.setFriends(friendMap.getOrDefault(client.getClientCode(), Collections.emptyList()));
            return result;
        }).collect(Collectors.toList());
    }

    // 模拟获取批量客户端编码的方法（根据你的业务实现）
    private List<String> getBatchClientCodes() {
        // 示例：返回5000个客户端编码
        return IntStream.rangeClosed(1, 5000)
                .mapToObj(i -> "CLIENT_" + i)
                .collect(Collectors.toList());
    }
}

4. 关键注意事项 #

• 事务兼容性：确保你的EntityManagerFactory和PlatformTransactionManager都使用AbstractRoutingDataSource，这样事务会自动绑定当前切换后的数据源。
• 线程安全：必须在finally块中恢复原数据源，避免Spring Batch线程池中的线程复用导致后续任务使用错误数据源。
• 性能优化：一次性查询5000条数据时，确保Postgres的Friend表在clientCode字段上有索引，避免全表扫描。

内容的提问来源于stack exchange，提问作者toch