API의 쿼리 개수 세기 - (2) JDBC, Spring AOP, Dynamic Proxy를 활용한 카운팅

지난 시간에 이어 이번에는 하이버네이트를 사용하지 않을 때 쿼리 개수를 카운팅하는 방법에 대해 알아보겠습니다. 하이버네이트를 사용할 때와 달리 특정 기술에 종속적이지 않고, JDBC가 자바의 표준 데이터 접근 기술이기 때문에 JDBC를 쓰든 MyBatis를 쓰든 JPA를 쓰든 적용할 수 있다는 특징이 있습니다. 이번에도 역시 아이디어는 같습니다. "쿼리가 실행될 때 query count 값을 1 증가시킨다."입니다. 하지만 이번에는 하이버네이트의 힘을 빌리지 않기 때문에, 쿼리가 실행되는 시점을 아는 것이 중요합니다. 어떤 방법으로 쿼리가 실행되었는지를 알 수 있을까요?

JDBC를 사용한 쿼리 카운팅 아이디어

자바는 데이터베이스의 종류에 상관없이 데이터베이스에 접속하고 쿼리를 실행하기 위해 JDBC API를 사용합니다. 그리고 JDBC API는 쿼리를 실행하기 위해 PreparedStatement 객체의 executeQuery 또는 execute나 executeUpdate를 사용합니다. 하이버네이트 등 ORM을 사용하더라도 기본적으로는 JDBC API를 사용하기 때문에 마찬가지입니다. 따라서, PreparedStatement.executeQuery, execute, executeUpdate의 실행 시점에 쿼리 카운터 객체의 카운트 값을 증가시키면 됩니다. 이렇게 쿼리를 실행하는 메서드가 세 개나 있지만, 앞으로는 편의상 executeQuery의 경우만 생각하도록 하겠습니다.

그렇다면 어떻게 해야 executeQuery의 실행 시점에 맞추어 카운트를 증가시키는 동작을 할 수 있을까요? 특정 메서드의 실행 시 카운트 값을 증가시키는 것은 부가 기능입니다. 따라서 AOP를 사용하면 된다는 생각을 해볼 수 있습니다.

하지만 문제가 있습니다. 스프링 AOP는 오직 스프링 빈에만 적용시킬 수 있습니다. 하지만 PreparedStatement 객체는 스프링 빈으로 등록할 수 없기 때문에 스프링 AOP를 적용할 수 없습니다. 그렇다면 PreparedStatement 객체를 만드는 Connection 객체는 어떨까요? Connection 역시 스프링 빈으로 만들 수 없습니다. Connection을 가져오는 DataSource까지 올라가야 스프링 AOP를 적용할 수 있는 것이죠. 때문에 Connection과 PreparedStatement를 다이나믹 프록시로 만들어서 이 문제를 해결하고, DataSource에 스프링 AOP를 적용할 수 있습니다. 다행히 Connection과 PreparedStatement 모두 인터페이스기 때문에 리플렉션을 사용하여 간단하게 프록시로 만들 수 있습니다. 최종 아이디어는 다음과 같은 그림이 됩니다.

여기서 잠깐, 꼭 다이나믹 프록시가 아니어도 괜찮습니다. 그냥 부가 기능을 담은 프록시 객체만 만들어 주면 됩니다. 하지만 그렇게 되면 Connection 클래스와 PreparedStatement에 정의된 다른 모든 메서드들을 모두 implement 해주어야 해서 불편하기 때문에 다이나믹 프록시로 만들기로 하겠습니다. 다이나믹 프록시를 만드는 방법에 대해서는 여기를 참고 해주세요.

쿼리 카운터는 지난 시간에 만들었던 ApiQueryCounter를 그대로 사용하도록 하겠습니다.

@Component
@RequestScope
@Getter
public class ApiQueryCounter {

    private int count;

    public void increaseCount() {
        count++;
    }
}

PreparedStatement의 다이나믹 프록시 만들기

먼저 쿼리 실행 시점에 카운트를 증가시키는 다이나믹 프록시부터 만들어보겠습니다. 원하는 동작을 설정할 InvocationHandler를 만들어야 합니다.

public class PreparedStatementProxyHandler implements InvocationHandler {

    private final Object preparedStatement;
    private final ApiQueryCounter apiQueryCounter;

    public PreparedStatementProxyHandler(final Object preparedStatement, final ApiQueryCounter apiQueryCounter) {
        this.preparedStatement = preparedStatement;
        this.apiQueryCounter = apiQueryCounter;
    }

    @Override // (1)
    public Object invoke(final Object proxy, final Method method, final Object[] args) throws Throwable {
        if (isExecuteQuery(method) && isInRequestScope()) { // (2)
            apiQueryCounter.increaseCount();
        }
        return method.invoke(preparedStatement, args); // (3)
    }

    private boolean isExecuteQuery(final Method method) {
        String methodName = method.getName();
        return methodName.equals("executeQuery") || methodName.equals("execute") || methodName.equals("executeUpdate");
    }

    private boolean isInRequestScope() {
        return Objects.nonNull(RequestContextHolder.getRequestAttributes());
    }
}

• (1) invoke 메서드는 다이나믹 프록시의 target의 메서드 호출을 가로챈다고 생각하면 됩니다. 이 메서드의 구현을 통해 메서드 내용을 확장할지, 어떤 방식으로 구현할지를 결정합니다.
• (2) invoke의 대상을 if문을 통해 지정해보도록 하겠습니다. 저희가 원하는 메서드는 executeQuery, execute, executeUpdate입니다. 또한 Request Scope의 빈인 ApiQueryCounter를 사용하기 때문에 이 요청이 API 요청 내에서 이루어지고 있는지 확인하는 분기도 추가해주겠습니다. 분기문을 통과하면 쿼리 카운터 빈의 카운트를 1 증가시킵니다.
• (3) 메서드의 실제 구현에는 손대지 않을 것이므로 원래 메서드를 실행하고 결과를 그대로 반환해줍니다.

이렇게 하면 쿼리 실행 메서드가 실행될 때 카운트를 증가시키는 동작을 할 수 있습니다.

Connection의 다이나믹 프록시 만들기

하지만 InvocationHandler를 만들기만 했지 실제로 사용되는 곳은 아직 없습니다. Connection이 prepareStatement 메서드를 호출할 때 PreparedStatement 대신 프록시를 반환하도록 해주어야 합니다. 이를 위해서는 Connection 역시도 다이나믹 프록시로 만들어줘야겠죠? 그를 위한 또다른 InvocationHandler를 만들어주도록 하겠습니다.

public class ConnectionProxyHandler implements InvocationHandler {

    private final Object connection;
    private final ApiQueryCounter apiQueryCounter;

    public ConnectionProxyHandler(final Object connection, final ApiQueryCounter apiQueryCounter) {
        this.connection = connection;
        this.apiQueryCounter = apiQueryCounter;
    }

    @Override
    public Object invoke(final Object proxy, final Method method, final Object[] args) throws Throwable {
        Object invokeResult = method.invoke(connection, args); // (1)
        if (method.getName().equals("prepareStatement")) {
            return Proxy.newProxyInstance(
                    invokeResult.getClass().getClassLoader(),
                    invokeResult.getClass().getInterfaces(),
                    new PreparedStatementProxyHandler(invokeResult, apiQueryCounter)
            );
        }
        return invokeResult;
    }
}

PreparedStatementProxyHandler와 내부 구현은 비슷합니다. prepareStatement라는 메서드에 분기 처리를 해주면 됩니다. 이 때 (1)을 조심해주셔야 하는데요, 다이나믹 프록시를 만들 때는 실제 메서드의 반환값, 그러니까 여기서는 PreparedStatement 객체가 필요합니다. 때문에 method.invoke를 먼저 실행해 주고, 다이나믹 프록시를 만들어야 하는 조건이면 해당 메서드의 반환값을 통해 다이나믹 프록시를 만들어주도록 하겠습니다. 이제 PreparedStatementProxyHandler를 InvocationHandler로 하는 Connection의 프록시는 prepareStatement 메서드를 실행하면 쿼리 카운팅이 가능한 프록시를 반환합니다.

AOP를 이용하여 Connection의 프록시 반환하기

이제 이 핸들러가 작동할 수 있는 Connection의 다이나믹 프록시를 만들어야 합니다. Connection 객체를 프록시로 대체하는 것은 DataSource.getConnection 메서드 호출 시 실행할 동작입니다. 앞서 설명했듯이 DataSource는 스프링 빈이므로, 스프링 AOP의 힘을 빌려보도록 하겠습니다.

@Component
@Aspect
public class ApiQueryCounterAop {

    private final ApiQueryCounter apiQueryCounter;

    public ApiQueryCounterAop(final ApiQueryCounter apiQueryCounter) {
        this.apiQueryCounter = apiQueryCounter;
    }

    @Around("execution(* javax.sql.DataSource.getConnection())")
    public Object getConnection(ProceedingJoinPoint proceedingJoinPoint) throws Throwable {
        Object connection = proceedingJoinPoint.proceed();
        return Proxy.newProxyInstance(
                connection.getClass().getClassLoader(),
                connection.getClass().getInterfaces(),
                new ConnectionProxyHandler(connection, apiQueryCounter)
        );
    }
}

AOP에 대한 자세한 설명은 우테코 테코블이나 토비의 스프링 같은 서적을 참고해주세요

DataSource.getConnection의 호출을 JoinPoint로 하고, 반환값이 Connection의 다이나믹 프록시가 되도록 AOP를 설정해주도록 하겠습니다. 이렇게 하면 DataSource로 부터 Connection을 받을 때 ConnectionProxyHandler에 작성한 대로 동작을 하는 프록시 객체를 받을 수 있습니다.

인터셉터에서 로그를 기록

이제 카운터를 다 만들었으니 로그로 기록할 차례입니다. 이전 게시글에서와 마찬가지로 LoggingInterceptor를 만들어주겠습니다.

@Slf4j
@Component
public class LoggingInterceptor implements HandlerInterceptor {

    private static final String QUERY_COUNT_LOG_FORMAT = "STATUS_CODE: {}, METHOD: {}, URL: {}, QUERY_COUNT: {}";
    private static final String QUERY_COUNT_WARNING_LOG_FORMAT = "쿼리가 {}번 이상 실행되었습니다.";

    private static final int QUERY_COUNT_WARNING_STANDARD = 10;

    private final ApiQueryCounter apiQueryCounter;

    public LoggingInterceptor(final ApiQueryCounter apiQueryCounter) {
        this.apiQueryCounter = apiQueryCounter;
    }

    @Override
    public void afterCompletion(final HttpServletRequest request, final HttpServletResponse response,
                                final Object handler, final Exception ex) {
        final int queryCount = apiQueryCounter.getCount();

        log.info(QUERY_COUNT_LOG_FORMAT, response.getStatus(), request.getMethod(), request.getRequestURI(),
                queryCount);
        if (queryCount >= QUERY_COUNT_WARNING_STANDARD) {
            log.warn(QUERY_COUNT_WARNING_LOG_FORMAT, QUERY_COUNT_WARNING_STANDARD);
        }
    }
}

개수에 맞게 정상적으로 카운트가 계산되는 것을 확인할 수 있었습니다!