[Spring Batch] Spring Batch, QueryDSL 적용기

🎯 Spring Batch 에서 QueryDSL 를 사용하는게 좋을까?

---

Spring Batch 에서는 데이터를 읽을때(reader) 다양한 ItemReader 를 제공합니다. (ex JdbcCursorItemReader, JdbcPagingItemReader 등등..)

 

그렇다면 QueryDSL 을 이용한 ItemReader 도 지원하지 않을까해서 찾아봤지만,, 아쉽게도 지원하지 않습니다. 그래서 지금껏 Spring Batch 를 이용해 DB 에서 JPA 를 이용해 데이터를 조회할때 JpaPagingItemReader 를 사용해 쿼리를 직접 작성해줬습니다.

 

하지만, 직접 쿼리를 작성하는 만큼 타입 안정성, 자동완성, 컴파일 단계 문법체크 등의 이점을 가질 수 없었습니다. 간단한 쿼리라면 괜찮은데 보통 성능을 개선할때 쿼리를 튜닝하며 복잡해지는 경우가 많았습니다.

 

실제로 현재 프로젝트에서 사용중인 Batch 코드는 다음과 같습니다. 딱 봐도 정말 복잡한 쿼리죠 ? 그래서 항상 QueryDSL 를 이용한 ItemReader 가 있으면 정말 좋았을텐데.. 라는 생각을 저만 한게 아니였습니다.

 

그래서 이번 포스팅에서는 Spring Batch 에서 QueryDSL 를 사용하는 방법에 대해 설명드리려 합니다.

  public static JpaPagingItemReader<SettlementAggregation> settlementReader(EntityManagerFactory entityManagerFactory, DateParameter dateParameter) {
        String query = """
            SELECT new com.dto.SettlementAggregation(
                t.shopId,
                t.shopName,
                SUM(CASE WHEN t.status = 'CANCEL' THEN t.price ELSE 0 END),
                SUM(CASE WHEN t.status <> 'CANCEL' THEN t.price ELSE 0 END),
                SUM(CASE WHEN t.status <> 'CANCEL' THEN 
                    CASE 
                        WHEN t.discountType = 'VIP_DISCOUNT' THEN (t.price * 0.9 - 1000)
                        WHEN t.discountType = 'FIRST_ORDER_DISCOUNT' THEN (t.price * 0.95 - 1000)
                        ELSE (t.price - 1000)  
                    END
                ELSE 0 END),
                   MAX(t.completionDateTime)
            )
            FROM OrderTransaction t
            WHERE FUNCTION('DATE', t.completionDateTime) = :requestDate
            GROUP BY t.shopId
        """;

        return new JpaPagingItemReaderBuilder<SettlementAggregation>()
                .name("settlementReader")
                .entityManagerFactory(entityManagerFactory)
                .pageSize(100)
                .queryString(query)
                .parameterValues(Map.of("requestDate", dateParameter.getRequestDate()))
                .build();
    }

 

🎯 AbstractPagingItemReader 를 상속받아 구현하기

---

✅ AbstractPagingItemReader 란?

 

AbstractPagingItemReader는 Spring Batch에서 페이징(Paging) 기반으로 데이터를 읽어오는 데 사용하는 추상 클래스입니다.
이를 확장(extends)하면 데이터베이스에서 페이징을 이용하여 데이터를 가져오는 커스텀 ItemReader를 만들 수 있습니다.

 

AbstractPagingItemReader 에서 doReadPage 는 JPQL 를 만들고 실행시킵니다. 즉, doReaderPage 에서 JPQL 를 만들고 실행시키는 로직을 변경한다면 충분히 QueryDSL 를 이용해 ItemReader 를 구현할 수 있습니다.

 

하지만, createQuery 메서드의 접근 제어자가 private 로 지정되어 있어 해당 부분만 오버라이드할 수 없어 AbstractPagingItemReader 코드를 복사해 직접 구현체를 만들어줘야 합니다.

@Override
@SuppressWarnings("unchecked")
protected void doReadPage() {

   EntityTransaction tx = null;

   if (transacted) {
      tx = entityManager.getTransaction();
      tx.begin();

      entityManager.flush();
      entityManager.clear();
   } // end if

   Query query = createQuery().setFirstResult(getPage() * getPageSize()).setMaxResults(getPageSize());

   if (parameterValues != null) {
      for (Map.Entry<String, Object> me : parameterValues.entrySet()) {
         query.setParameter(me.getKey(), me.getValue());
      }
   }

 

우리는 이를 활용해 커스텀 QuerydslPagingItemReader 를 만들 수 있습니다.

 

✅ QuerydslPagingItemReader 만들어보기

 

먼저 AbstractPagingItemReader 를 상속받아 구현한 QuerydslPagingItemReader 입니다.

public class QuerydslPagingItemReader<T> extends AbstractPagingItemReader<T> {

    protected final Map<String, Object> jpaPropertyMap = new HashMap<>();
    protected EntityManagerFactory entityManagerFactory;
    protected EntityManager entityManager;
    protected Function<JPAQueryFactory, JPAQuery<T>> queryFunction;
    protected boolean transacted = true;//default value

    protected QuerydslPagingItemReader() {
        setName(ClassUtils.getShortName(QuerydslPagingItemReader.class));
    }

    public QuerydslPagingItemReader(EntityManagerFactory entityManagerFactory,
                                    int pageSize,
                                    Function<JPAQueryFactory, JPAQuery<T>> queryFunction) {
        this();
        this.entityManagerFactory = entityManagerFactory;
        this.queryFunction = queryFunction;
        setPageSize(pageSize);
    }

    public void setTransacted(boolean transacted) {
        this.transacted = transacted;
    }

    @Override
    protected void doOpen() throws Exception {
        super.doOpen();

        entityManager = entityManagerFactory.createEntityManager(jpaPropertyMap);
        if (entityManager == null) {
            throw new DataAccessResourceFailureException("Unable to obtain an EntityManager");
        }
    }

    @Override
    @SuppressWarnings("unchecked")
    protected void doReadPage() {

        clearIfTransacted();

        JPAQuery<T> query = createQuery()
                .offset(getPage() * getPageSize())
                .limit(getPageSize());

        initResults();

        fetchQuery(query);
    }

    protected void clearIfTransacted() {
        if (transacted) {
            entityManager.clear();
        }
    }

    protected JPAQuery<T> createQuery() {
        JPAQueryFactory queryFactory = new JPAQueryFactory(entityManager);
        return queryFunction.apply(queryFactory);
    }

    protected void initResults() {
        if (CollectionUtils.isEmpty(results)) {
            results = new CopyOnWriteArrayList<>();
        } else {
            results.clear();
        }
    }

    protected void fetchQuery(JPAQuery<T> query) {
        if (!transacted) {
            List<T> queryResult = query.fetch();
            for (T entity : queryResult) {
                entityManager.detach(entity);
                results.add(entity);
            }
        } else {
            results.addAll(query.fetch());
        }
    }
    
    @Override
    protected void doClose() throws Exception {
        entityManager.close();
        super.doClose();
    }
}

 

 

 기존의 JpaPagingItemReader 를 참고해 새로운 구현체를 직접 만들었습니다. 여기서 주의깊게 봐야할 점은 생성자에 Function<JPAQueryFactory, JPAQuery<T> queryFunction> 를 주입받은 점 입니다. 이는 QueryDSL 를 람다식 으로 받기 위함 입니다. 즉, 우리는 Spring Batch Job 을 만들때 ItemReader 를 만드는 과정에서 쿼리가 아닌 QueryDSL 를 이용한 람다식을 만들어 주입해 줄 수 있습니다.

 

 

✅ QuerydslPagingItemReader 테스트

 

자! 이제 만든 커스텀 Reader 가 정상적으로 작동하는지 테스트 해보겠습니다.

@SpringBatchTest
@ActiveProfiles("test")
@ExtendWith(SpringExtension.class)
@SpringBootTest(classes = {TestBatchLegacyConfig.class, QueryDslConfig.class})
public class QuerydslPagingItemReaderTest {

    @Autowired
    private EntityManagerFactory emf; // EntityManagerFactory 주입

    @Autowired
    private UserRepository userRepository;

    @Test
    public void 커스텀_reader_테스트() throws Exception {
        // given
        User savedUser = userRepository.save(createUser());
        User savedUser2 = userRepository.save(createUser());

        QUser qUser = QUser.user;
        int pageSize = 1;

        // when
        QuerydslPagingItemReader<User> reader = new QuerydslPagingItemReader<>(
                emf,
                pageSize,
                queryFactory -> queryFactory.selectFrom(qUser)
        );

        reader.open(new ExecutionContext());
        User readUser = reader.read(); // 첫 번째 사용자 읽기
        User readUser2 = reader.read(); // 두 번째 사용자 읽기
        reader.close(); // 리더 종료

        // then
        assertThat(readUser).isNotNull();
        assertThat(readUser.getId()).isEqualTo(savedUser.getId());
        assertThat(readUser.getEmail()).isEqualTo(savedUser.getEmail());
        assertThat(readUser2.getId()).isEqualTo(savedUser2.getId());
        assertThat(readUser2.getEmail()).isEqualTo(savedUser2.getEmail());
    }
}

 

 

QuerydslPagingItemReader를 사용하기 위해서는 EntityManagerFactory 를 주입해야 합니다. 저같은 경우 QueryDslConfig 에서 EntityManagerFactory 를 생성하고 있습니다.

 

또한 ItemReader 를 단독으로 테스트 하기 위해서는 별도의 실행환경 (ExecutionContext)을 등록 해줘야만 합니다. 즉, open() 을 하지 않으면 EntityManagerFactory 가 등록되지 않습니다. 

 

 

자! 이제 QuerydslPagingItemReader 가 정상적으로 데이터를 읽어 들어옴을 확인했습니다. 그리고 pageSize 를 1 로 설정하고 총 2개의 데이터를 읽어들어올때 총 3번의 페이징 조회가 실행됩니다. 두번이 아닌, 세번의 페이징 조회가 실행된 이유는 마지막 쿼리에서 데이터가 남아 있는지 확인하기 때문입니다.

 

자! 이제 만든 QuerydslPagingItemReader 를 적용해보겠습니다.

 

✅ QuerydslPagingItemReader 적용 전 코드 

    public static JpaPagingItemReader<SettlementAggregation> settlementReader(EntityManagerFactory entityManagerFactory, DateParameter dateParameter) {
        String query = """
            SELECT new com.dto.SettlementAggregation(
                t.shopId,
                t.shopName,
                SUM(CASE WHEN t.status = 'CANCEL' THEN t.price ELSE 0 END),
                SUM(CASE WHEN t.status <> 'CANCEL' THEN t.price ELSE 0 END),
                SUM(CASE WHEN t.status <> 'CANCEL' THEN 
                    CASE 
                        WHEN t.discountType = 'VIP_DISCOUNT' THEN (t.price * 0.9 - 1000)
                        WHEN t.discountType = 'FIRST_ORDER_DISCOUNT' THEN (t.price * 0.95 - 1000)
                        ELSE (t.price - 1000)  
                    END
                ELSE 0 END),
                   MAX(t.completionDateTime)
            )
            FROM OrderTransaction t
            WHERE FUNCTION('DATE', t.completionDateTime) = :requestDate
            GROUP BY t.shopId
        """;

        return new JpaPagingItemReaderBuilder<SettlementAggregation>()
                .name("settlementReader")
                .entityManagerFactory(entityManagerFactory)
                .pageSize(100)
                .queryString(query)
                .parameterValues(Map.of("requestDate", dateParameter.getRequestDate()))
                .build();
    }

 

현재 JpaPagingItemReader 를 이용해 DB 에서 데이터를 조회하고 있습니다. 직접 쿼리를 만들어 주입하므로 타입 안정성, 자동완성, 컴파일 단계 문법체크 를 할 수 없는 상황입니다. 

 

QuerydslPagingItemReader 적용 코드

public static ItemReader<SettlementAggregation> settlementReader(EntityManagerFactory entityManagerFactory, DateParameter requestDate) {
    QOrderTransaction t = QOrderTransaction.orderTransaction;

    NumberExpression<Double> cancelPriceSum = new CaseBuilder()
            .when(t.status.eq(OrderStatus.CANCEL)).then(t.price)
            .otherwise(0.0)
            .sum();

    NumberExpression<Double> nonCancelPriceSum = new CaseBuilder()
            .when(t.status.ne(OrderStatus.CANCEL)).then(t.price)
            .otherwise(0.0)
            .sum();

    NumberExpression<Double> discountedPriceSum = new CaseBuilder()
            .when(t.status.ne(OrderStatus.CANCEL))
            .then(new CaseBuilder()
                    .when(t.discountType.eq(DiscountType.VIP_DISCOUNT))
                    .then(t.price.multiply(0.9).subtract(1000))
                    .when(t.discountType.eq(DiscountType.FIRST_ORDER_DISCOUNT))
                    .then(t.price.multiply(0.95).subtract(1000))
                    .otherwise(t.price.subtract(1000)))
            .otherwise(0.0)
            .sum();

    Function<JPAQueryFactory, JPAQuery<SettlementAggregation>> queryFunction = queryFactory ->
            queryFactory.select(Projections.constructor(SettlementAggregation.class,
                            t.shopId,
                            t.shopName,
                            cancelPriceSum,
                            nonCancelPriceSum,
                            discountedPriceSum,
                            t.completionDateTime.max()
                    ))
                    .from(t)
                    .where(t.completionDateTime.year().eq(requestDate.getRequestDate().getYear())
                            .and(t.completionDateTime.month().eq(requestDate.getRequestDate().getMonthValue()))
                            .and(t.completionDateTime.dayOfMonth().eq(requestDate.getRequestDate().getDayOfMonth())))
                    .groupBy(t.shopId);

    return new QuerydslPagingItemReader<>(entityManagerFactory, 100, queryFunction);
}

 

이처럼 QueryDsl 을 이용해 ItemReader 를 구현하면, 컴파일시 타입에 대한 오류를 잡아줄 수 있습니다. 또한 동적 쿼리 처리가 보다 간편해졌고 JPQL 과 드리게 DB에 종속적인 함수를 사용하지 않아 DB 변경시 호환성이 더 좋아졌습니다. 꼭 QueryDsl 를 이용한 reader 방식이 좋다고 말할 순 없지만 복잡한 동적 쿼리가 필요하거나, 유지보수성이 중요한경우 이러한 방식이 더 나은 선택이라 생각했습니다.

 

 

참고 글

https://techblog.woowahan.com/2662/

Spring Batch와 Querydsl | 우아한형제들 기술블로그