다이내믹 프록시를 활용한 JPA QueryCounter 구현기

들어가며

JPA와 같은 ORM을 사용할 때, 지연 로딩과 같은 기능을 무분별하게 이용하다보면 원래 의도했던 개수보다 더 많은 쿼리가 발생하기도 한다. ORM에서 흔히 발생하는 N + 1 문제도 개발자의 의도와 다르게 N번의 쿼리가 더 발생하게 된다.

이렇듯 ORM을 사용하다보면 예상하지 못한 개수의 쿼리가 발생하기도 하는데, 이를 한눈에 파악하기는 쉽지 않다. 물론 JPA의 show-sql 기능을 사용하면 로그에 쿼리가 출력되기는 하지만, 좀 더 쉽게 알아볼 수 있는 방법이 필요하다는 생각이 들었고, 이를 위해 다이내믹 프록시를 활용하여 JPA QueryCounter를 구현해보기로 하였다.

여기서는 Spring Data JPA를 기반으로 설명한다는 점을 알고가도록 하자.

기본 원리

JPA QueryCounter를 구현하기 위해 활용할 부분은 JPA가 JDBC API를 이용한다는 점이다.

예를 들어, 새로운 데이터를 삽입한다고 하면, 개발자는 객체를 마치 컬렉션에 저장하듯 JPA에 저장하면 JPA가 내부적으로 엔티티를 분석하고, 쿼리를 생성한 뒤, JDBC API를 이용하여 데이터베이스에 객체를 저장해준다.

즉, 우리가 흔히 사용하는 다음과 같은 코드는

postRepository.save(new Post(...));

내부적으로 아래와 같은 과정을 가진다는 것이다.

Connection conn = dataSource.getConnection();
PreparedStatement pstmt = conn.prepareStatement("INSERT INTO POST ...");

...

pstmt.executeUpdate();

이를 통해 JPA를 사용하여 데이터베이스에 접근한다는 것은 결국 Connection 객체를 통해 PrepareStatement 객체를 생성하여 쿼리를 실행하는 것이라는 것을 알 수 있다. 그렇다면 Connection 객체의 prepareStatement 메서드가 호출되는 횟수를 세어준다면 쿼리가 발생한 횟수를 계산할 수 있지 않을까?

바로 이 부분이 이번 포스트에서 구현하려는 QueryCounter의 기본 원리이다.

그렇다면 어떻게 Connection 객체의 prepareStatement 메서드가 호출되는 횟수를 셀 수 있을까?

이는 이전 포스팅에서 다루어보았던 다이내믹 프록시를 이용하면 된다.

다이내믹 프록시 적용

구체적으로 이야기해보자면, Connection에 대한 다이내믹 프록시를 생성하고, 다이내믹 프록시가 타깃 Connection 객체에 작업을 위임할 때 해당 작업이 prepareStatement 메서드라면 개수를 세어주는 부가 작업을 수행해줌으로써 쿼리의 개수를 카운트하는 것이다.

따라서 타깃 Connection 객체로 PrepareStatement 생성 작업을 위임할 때 개수를 세어주는 부가 작업을 수행하는 InvocationHandler를 구현해 줄 필요가 있다.

그전에 먼저, 쿼리의 개수를 저장하는 Counter객체를 구현해보도록 하자.

public class Counter {

        private int count;
        private boolean flag;

        public Counter() {
        }

        public void startQueryCount() {
                this.flag = true;
        }

        public void countQuery() {
                this.count += 1;
        }

        public int getQueryCount() {
                return this.count;
        }

        public boolean getFlag() {
                return this.flag;
        }

        public void printQueryCount() {
                System.out.println("ㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡ");
                System.out.println("Query Count : " + this.count + "번");
                System.out.println("ㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡ");
                
                clearQueryCount();
        }

        public void clearQueryCount() {
                this.flag = false;
                this.count = 0;
        }

}

startQueryCount 메서드를 호출한 시점부터 발생한 쿼리들의 개수를 세어주도록 구현하였다. flag 값을 두지 않으면 DDL 쿼리들도 카운팅되기 때문에 정확한 개수를 파악할 수 없으므로, flag 값을 기반으로 쿼리의 개수를 세는 기점을 알려줄 필요가 있다.

다음은 앞서 구현한 Counter 객체를 사용하여 쿼리의 개수를 세어주는 부가 작업을 하는 QueryCountHandler를 구현해보도록 하자.

public class QueryCountHandler implements InvocationHandler {

        private Object target;
        private Counter counter;

        public QueryCountHandler(Object target, Counter counter) {
                this.target = target;
                this.counter = counter;
        }

        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
                try {
                        Object ret = method.invoke(target, args);
                        if (counter.getFlag() && method.getName().startsWith("prepareStatement")) {
                                counter.countQuery();
                        }
                        return ret;
                } catch (InvocationTargetException e) {
                    	throw e.getTargetException();
                }
        }
}

QueryCountHandler의 기능은 타깃 Connection 객체로 작업을 위임한 뒤, 해당 작업이 prepareStatement 메서드라면 Counter 객체의 countQuery 메서드를 통해 개수를 세어주도록 하는 것이다.

이러한 QueryCountHandler를 통해 비즈니스 로직에서 발생하는 쿼리의 횟수를 구할 수 있게 된다.

이제 이렇게 구현한 QueryCountHandler를 가지고 Connection을 위한 다이내믹 프록시를 생성해서 적용해주어야 하는데, 어디에 적용해주어야 할지에 대해서 생각해보도록 하자.

바로 Connection을 생성해주는 DataSource이다. 우리는 실행되는 쿼리의 횟수를 세어주기 위해서 Connection 다이내믹 프록시를 사용하여야 하므로, DataSource에서 반환해주는 Connection이 앞서 구현한 다이내믹 프록시가 되도록 해야할 것이다.

Connection conn = dataSource.getConnection();	// 이 Connection이 다이내믹 프록시이면 된다.

따라서 DataSource를 상속하는 QueryCountDataSource를 구현해서 Connection을 반환해주는 getConnection 메서드가 다이내믹 프록시를 반환하도록 해보자.

public class QueryCountDataSource implements DataSource {

        private DataSource dataSource;
        private Counter counter;

        public QueryCountDataSource(DataSource dataSource, Counter counter) {
                this.dataSource = dataSource;
                this.counter = counter;
        }

        @Override
        public Connection getConnection() throws SQLException {
                Connection conn = dataSource.getConnection();
                return (Connection) Proxy.newProxyInstance(
                            getClass().getClassLoader(),
                            new Class[] { Connection.class },
                            new QueryCountHandler(conn, counter)
                );
        }

    	...
}

위 코드를 보면, 타깃 Connection 객체를 생성해서 이를 다이내믹 프록시로 감싸서 반환하는 것을 알 수 있다.

이렇게 다이내믹 프록시를 반환하는 QueryCountDataSource까지 구현해보았는데, 이제는 JPA가 이 다이내믹 프록시를 사용하도록 해야 한다. 따라서 이를 유도해보도록 하자.

EntityManagerFactory 커스텀

JPA는 EntityManager를 통해 DB 접근 작업을 수행한다. Spring Data JPA에서 흔히 사용하는 JpaRepository도 이를 구현한 SimpleJpaRepository를 보면 내부적으로 EntityManager를 사용한다. 그리고 이러한 EntityManager를 사용하여 DB 접근 작업을 수행하면 앞서 언급한 바와 같이 결국 JDBC API를 이용하게 된다.

그렇다면 우리는 EntityManager가 사용하는 Connection을 앞서 구현한 다이내믹 프록시로 바꿔주기만 한다면 동일한 기능을 제공하면서 쿼리의 개수를 셀 수 있는 목표에 도달할 수 있을 것이다. 이를 위해서는 EntityManager를 생성해주는 EntityManagerFactory를 커스텀해야 한다.

EntityManagerFactory를 커스텀하는 방법으로는 LocalContainerEntityManagerFactoryBean을 스프링 빈으로 등록하는 방법이 있다. 다음 코드를 살펴보도록 하자.

LocalContainerEntityManagerFactoryBean은 JPA EntityManagerFactory를 생성해주는 팩토리 빈이다.

@Configuration
public class QueryCountConfig {

        @Bean
        public DataSource dataSource() {
                return DataSourceBuilder.create()
                        .driverClassName("org.h2.Driver")
                        .url("jdbc:h2:mem:testdb")
                        .username("sa")
                        .password("")
                        .build();
        }

        @Bean
        public Counter counter() {
                return new Counter();
        }

        @Bean
        public DataSource queryCountDataSource() {
                return new QueryCountDataSource(dataSource(), counter());
        }

        @Bean
        public LocalContainerEntityManagerFactoryBean entityManagerFactory() {
                LocalContainerEntityManagerFactoryBean em = new LocalContainerEntityManagerFactoryBean();
                em.setDataSource(queryCountDataSource());	// 다이내믹 프록시를 위한 설정
                em.setPackagesToScan("com.example.querycounter.domain");

                HibernateJpaVendorAdapter vendorAdapter = new HibernateJpaVendorAdapter();
                vendorAdapter.setShowSql(true);
                vendorAdapter.setGenerateDdl(true);

                em.setJpaVendorAdapter(vendorAdapter);

                return em;
        }

        ...
}

위 코드에서 entityManagerFactory 메서드를 보면 LocalContainerEntityManagerFactoryBean를 생성해서 DataSource를 앞서 구현한 QueryCountDataSource로 설정해주어서 DataSource에서 반환해주는 Connection이 앞서 구현한 다이내믹 프록시가 되도록 하였다.

이렇게 커스텀한 EntityManagerFactory를 스프링 빈으로 등록하면 이제는 EntityManagerFactory가 생성하는 EntityManager가 내부적으로 다이내믹 프록시를 사용하여 DB 접근 작업을 수행하게 된다. 그리고 이를 사용하는 JpaRepository에서 발생하는 쿼리의 개수를 셀 수 있게 된다.

QueryCounter 동작방식

여기까지 구현한 QueryCounter의 동작방식을 그림으로 나타내보면 위와 같다.

결국 핵심은 JPA가 내부적으로 이용하는 JDBC API에서 Connection객체를 다이내믹 프록시로 감싸고, Connection 객체가 PrepareStatement 객체를 생성할 때마다 카운팅해주는 것이다.

그리고 이러한 다이내믹 프록시를 JPA EntityManager가 이용하도록 하면 QueryCounter를 구현할 수 있게 된다.

QueryCounter 기능 테스트

이제 QueryCounter가 올바르게 동작하는지 알아보도록 하자.

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
public class QueryCounterTest {

        @Autowired
        private Counter counter;

        @Autowired
        private UserRepository userRepository;

        @Autowired
        private PostRepository postRepository;

        @Test
        void queryCounterTest() {
                counter.startQueryCount();

                User user = userRepository.save(new User("yeongkee"));

                Post post = Post.builder()
                            .writer(user)
                            .contents("post contents")
                            .build();
                postRepository.save(post);

                assertThat(counter.getQueryCount()).isEqualTo(2);
                counter.printQueryCount();
        }

}

위 테스트 코드에서 2번 save 메서드를 호출하였고, 정상적으로 2번 카운팅된 것을 확인할 수 있다.

다음은 N + 1 문제를 유도해보도록 하자.

@Transactional
@Service
public class UserService {

        private UserRepository userRepository;

        public UserService(UserRepository userRepository) {
                this.userRepository = userRepository;
        }

        public List<String> getAllPostContents() {
                return userRepository.findAll()
                            .stream()
                            .map(User::getPosts)
                            .flatMap(Collection::stream)
                            .map(Post::getContents)
                            .collect(Collectors.toList());
        }

}

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
public class QueryCounterTest {

        @Autowired
        private Counter counter;

        @Autowired
        private UserService userService;

        @Autowired
        private UserRepository userRepository;

        @Autowired
        private PostRepository postRepository;

        @BeforeEach
        void setUp() {
                for (int i = 1; i <= 10; i++) {
                        User user = userRepository.save(new User(i + "번 user"));

                        Post post = Post.builder()
                                    .writer(user)
                                    .contents(i + "번 post")
                                    .build();

                        postRepository.save(post);
                }
        }

        @Test
        void N_Plus_1_Problem() {
                counter.startQueryCount();

                List<String> userPostContentsList = userService.getAllPostContents();

                assertThat(userPostContentsList).hasSize(10);
                assertThat(counter.getQueryCount()).isEqualTo(11);
                counter.printQueryCount();
        }

}

이번 테스트 코드는 사용자 10명이 각각 1개의 포스트를 등록한 상황에서 모든 사용자 포스트의 Contents를 가져오는 것을 테스트이다.

UserService의 getAllPostContents 메서드의 로직은 모든 사용자를 조회하는 쿼리가 1번 발생하고, 조회된 사용자의 수 N 만큼 지연로딩으로 설정된 포스트를 조회하는 쿼리가 또 발생하게 되는 로직이다.

그렇다면 10명의 사용자가 각각 1개의 포스트를 등록했으므로, 1 + 10(N) = 11번의 쿼리가 발생할 것이다. QueryCounter를 통해 쿼리 개수를 확인해보도록 하자.

총 11번의 쿼리가 발생하고, 정상적으로 카운팅된 것을 확인할 수 있다.

마치며

JPA를 사용할 때 의도와는 다른 개수의 쿼리가 발생하기도 하고, 이를 콘솔에 출력된 쿼리로 매번 파악하기가 번거로워 QueryCounter 구현 시작하였는데, 생각보다 더 쉽지 않아서 오래걸렸던 것 같다. 스프링, JPA, JDBC 등 여러 개의 개념을 조합해야 하다보니 동작방식(Fig 2)을 이해하고, 완성하기까지 정말 많이 지우고 그리고를 반복했던 것 같다. 하지만 많은 실패를 겪었던 만큼 스프링과 JPA, 다이내믹 프록시에 좀 더 가까워질 수 있었던 시간이 되었던 것 같다.

코드는 QueryCounter 코드에서 확인하실 수 있습니다.
혹시 잘못된 부분이 있다면 지적 혹은 조언 부탁드립니다:)

'Backend > Spring' 카테고리의 다른 글

테스트 격리하기 (0)	2022.05.04
트랜잭션 전파 알아보기 (0)	2022.04.26
AOP와 빈 후처리기를 이용한 부가 기능 분리 (0)	2022.03.26
다이내믹 프록시를 이용한 부가 기능 분리 (0)	2022.03.12
의존관계 주입(DI)과 객체 지향 설계 (0)	2022.03.02