Spring Batch와 Criteria API의 만남(2)

 

 

이전 포스트에서는 Spring Batch Reader에서 Criteria API를 활용하기 위해 Custom Reader(= CriteriaPagingItemReader) 를 만드는 방법과 적용 방식을 정리했습니다. 이번 글에서는 그 Reader를 실제 배치에 적용하며 만났던 문제와, 이를 어떻게 개선했는지에 대해 다룹니다.

 

---

CriteriaPagingItemReader에서 발생한 타임아웃 이슈

CriteriaPagingItemReader를 여러 배치에 적용하던 중, 대용량 마이그레이션에서 DB 타임아웃 문제가 발생했습니다. 페이지 수가 커질수록 조회 시간이 길어지고, 결국 커넥션 타임아웃으로 배치가 중단되는 현상이었습니다.

원인은 페이징의 구조적 한계였습니다. JPA가 페이지가 증가할 때마다 offset(= setFirstResult)이 계속 커지면서, DB가 앞쪽 레코드를 “건너뛰는” 비용이 누적되어 지연이 눈덩이처럼 커졌습니다. 이 지점부터는 단순 페이징으로 버티기 어렵다고 판단했고, Cursor 기반으로 Criteria API를 사용하는 방향을 검토하게 되었습니다.

 

---

Criteria API를 조금 더 범용적으로 쓰게 만들 수 없을까?

처음 떠오른 아이디어는 당연히 CriteriaCursorItemReader 를 새로 만드는 것이었습니다. 하지만 기능 중복과 유지보수 비용을 고려하면 비효율적이라 판단했습니다. 대신 이미 존재하는 컴포넌트를 조합하는 쪽, 즉 JpaQueryProvider를 커스터마이징해 Criteria API를 지원하도록 만들면, JpaPagingItemReader와 JpaCursorItemReader 양쪽에서 동일한 Criteria 정의를 재사용할 수 있어 확장성과 일관성이 좋아집니다.

 

---

 

JpaCriteriaQueryProvider 의 탄생

아래는 해당 아이디어를 구현한 JpaCriteriaQueryProvider 입니다.

public class JpaCriteriaQueryProvider<T> extends AbstractJpaQueryProvider {

    private Function<EntityManager, CriteriaQuery<T>> criteriaQueryGenerator;

    @Override
    public Query createQuery() {
        EntityManager entityManager = getEntityManager();
        CriteriaQuery<T> criteriaQuery = criteriaQueryGenerator.apply(entityManager);
        return entityManager.createQuery(criteriaQuery);
    }

    public void setCriteriaQueryGenerator(Function<EntityManager, CriteriaQuery<T>> criteriaQueryGenerator) {
        this.criteriaQueryGenerator = criteriaQueryGenerator;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        Assert.state(criteriaQueryGenerator != null, "CriteriaQueryGenerator must not be null");
    }
}

 

여기서 핵심은 criteriaQueryGenerator 입니다. 단순히 CriteriaQuery 인스턴스를 외부에서 주입하는 대신, Function<EntityManager, CriteriaQuery<T>> 형태로 생성 로직을 받습니다. 그 이유는 EntityManager의 생명주기와 관련이 있습니다.

• JpaItemReader 계열은 내부적으로 doOpen() 시점에 자체 EntityManager 를 엽니다.
• Cursor 기반 리더(JpaCursorItemReader)는 이 같은 EntityManager 로 생성된 Query를 사용해야 스트리밍/커서가 올바르게 유지됩니다.
• 다른 시점(혹은 다른 EntityManager)으로 만든 CriteriaQuery를 재사용하면, 트랜잭션/세션 경계가 어긋나 세션 종료 예외, 커서 닫힘, 플러시/캐시 불일치 등의 문제가 생길 수 있습니다.

그래서 createQuery()에서 getEntityManager()로 리더가 연 동일한 EntityManager를 가져온 뒤, 그걸 criteriaQueryGenerator.apply(entityManager)에 전달해 그 자리에서 CriteriaQuery를 생성하도록 설계했습니다.

 

---

적용 예시: JpaCursorItemReader + JpaCriteriaQueryProvider

 

이렇게 탄생한 JpaCriteriaQueryProvider는 아래의 방식처럼 적용합니다.

    @Bean
    @StepScope
    public JpaCursorItemReader<String> CustomReader() {
        JpaCursorItemReader<String> reader = new JpaCursorItemReader<>();
        reader.setName("customReader");
        reader.setEntityManagerFactory(entityManagerFactory);
        reader.setQueryProvider(createCriteriaQueryProvider());
        return reader;
    }

    private CriteriaQueryProvider<String> createCriteriaQueryProvider() {
        CriteriaQueryProvider<String> queryProvider = new CriteriaQueryProvider<>();
        queryProvider.setCriteriaQueryGenerator(this::createCriteriaQuery);
        return queryProvider;
    }

    private CriteriaQuery<String> createCriteriaQuery(EntityManager entityManager) {
        BasePeriod basePeriod = jobParameter.getPeriod();
        CriteriaBuilder cb = entityManager.getCriteriaBuilder();
        CriteriaQuery<String> query = cb.createQuery(String.class);

        Root<Entity> root = query.from(Entity.class);
        query.select(root.get(Entity_.ID));
        query.where(cb.between(root.get(Entity_.CREATED_DATE_TIME),
                getStartDateTime(), getEndDateTime()));
        query.orderBy(cb.desc(root.get(Entity_.ID)));

        return query;
    }

정리하면,

1. EntityManager를 인자로 받아 CriteriaQuery를 만드는 메서드를 준비하고,
2. 그 메서드를 JpaCriteriaQueryProvider의 criteriaQueryGenerator에 설정한 뒤,
3. 해당 Provider를 원하는 JpaItemReader의 queryProvider에 주입하면 됩니다.

이렇게 하면 Paging 리더와 Cursor 리더에서 동일한 Criteria 정의를 손쉽게 재사용할 수 있습니다.

 

---

 

 

마치며

여러 문제점을 거치며 최종적으로 JpaCriteriaQueryProvider를 도입했고, 기존의 CriteriaPagingItemReader 기반 코드를 걷어내어 QueryProvider 중심으로 단순화할 수 있었습니다. 그 결과 프로젝트 전반에서 PagingItemReader와 CursorItemReader 양쪽 모두에 Criteria API를 일관되게 적용할 수 있는 구조를 갖추게 되었습니다.

페이징의 offset 병목으로 인한 타임아웃 이슈를 만난다면, Cursor + QueryProvider + Criteria 조합을 한 번 고려해 보시길 추천드립니다.