append custom annotation & remove unused thread local object

Author: Hyeon9mak

build.gradle

@@ -8,7 +8,7 @@ plugins {
 }
 
 group = "hyeon9mak"
-version = "1.1.5-spring-boot-3"
+version = "2.0.1-spring-boot-3"
 
 java {
     toolchain {

src/main/java/hyeon9mak/multidatasourcequerycounter/ConnectionQueryMonitor.java

@@ -8,11 +8,11 @@ public class ConnectionQueryMonitor implements MethodInterceptor {
 
     private static final String JDBC_PREPARE_STATEMENT_METHOD_NAME = "prepareStatement";
 
-    private final QueryCounterRequestContextScopeHolder queryCounterRequestContextScopeHolder;
+    private final QueryCountPerRequest queryCountPerRequest;
     private final Object connection;
 
-    public ConnectionQueryMonitor(QueryCounterRequestContextScopeHolder queryCounterRequestContextScopeHolder, Object connection) {
-        this.queryCounterRequestContextScopeHolder = queryCounterRequestContextScopeHolder;
+    public ConnectionQueryMonitor(QueryCountPerRequest queryCountPerRequest, Object connection) {
+        this.queryCountPerRequest = queryCountPerRequest;
         this.connection = connection;
     }
 
@@ -22,7 +22,7 @@ public Object invoke(MethodInvocation invocation) throws Throwable {
 
         if (result != null && wasPreparedStatementInvoked(invocation)) {
             ProxyFactory proxyFactory = new ProxyFactory(result);
-            proxyFactory.addAdvice(new PreparedStatementQueryMonitor(queryCounterRequestContextScopeHolder));
+            proxyFactory.addAdvice(new PreparedStatementQueryMonitor(queryCountPerRequest));
             return proxyFactory.getProxy();
         }
 

src/main/java/hyeon9mak/multidatasourcequerycounter/CountQueries.java

@@ -0,0 +1,15 @@
+package hyeon9mak.multidatasourcequerycounter;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Target({ElementType.METHOD, ElementType.TYPE})
+@Retention(RetentionPolicy.RUNTIME)
+public @interface CountQueries {
+    /**
+     * prefix for query count log message.
+     */
+    String prefix() default "";
+}

src/main/java/hyeon9mak/multidatasourcequerycounter/HikariDataSourceQueryCounter.java

@@ -2,9 +2,9 @@
 
 import jakarta.servlet.http.HttpServletRequest;
 import org.aspectj.lang.ProceedingJoinPoint;
-import org.aspectj.lang.annotation.After;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
+import org.aspectj.lang.reflect.MethodSignature;
 import org.springframework.stereotype.Component;
 import org.springframework.web.context.request.RequestContextHolder;
 import org.springframework.web.context.request.ServletRequestAttributes;
@@ -24,21 +24,29 @@ public HikariDataSourceQueryCounter(QueryCountPerRequest queryCountPerRequest, Q
     @Around("execution( * com.zaxxer.hikari.HikariDataSource.getConnection())")
     public Object aroundConnection(ProceedingJoinPoint joinPoint) throws Throwable {
         Object connection = joinPoint.proceed();
-        QueryCounterRequestContextScopeHolder queryCounterRequestContextScopeHolder = new QueryCounterRequestContextScopeHolder(queryCountPerRequest);
-        ConnectionQueryMonitor connectionQueryMonitor = new ConnectionQueryMonitor(queryCounterRequestContextScopeHolder, connection);
+        ConnectionQueryMonitor connectionQueryMonitor = new ConnectionQueryMonitor(queryCountPerRequest, connection);
         return connectionQueryMonitor.getProxy();
     }
 
-    @After("within(@org.springframework.web.bind.annotation.RestController *)")
-    public void afterApiFinished() {
+    @Around("@annotation(countQueries)")
+    public Object aroundCountQueriesMethod(ProceedingJoinPoint joinPoint, CountQueries countQueries) throws Throwable {
+        Object result = joinPoint.proceed();
+
         ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
 
         if (isInRequestScope(attributes)) {
             HttpServletRequest request = attributes.getRequest();
-            queryCountPerRequest.updateApiUrl(request.getMethod() + request.getRequestURI());
+            queryCountPerRequest.updateApiUrl(request.getMethod() + " " + request.getRequestURI());
+        } else {
+            MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
+            String methodName = methodSignature.getMethod().getName();
+            String className = methodSignature.getDeclaringType().getSimpleName();
+            queryCountPerRequest.updateApiUrl(countQueries.prefix() + className + "." + methodName);
         }
 
         queryCountLogger.logQueryCount(queryCountPerRequest);
+
+        return result;
     }
 
     private boolean isInRequestScope(ServletRequestAttributes attributes) {

src/main/java/hyeon9mak/multidatasourcequerycounter/PreparedStatementQueryMonitor.java

@@ -11,10 +11,10 @@ public class PreparedStatementQueryMonitor implements MethodInterceptor {
 
     private static final List<String> QUERY_METHODS = Arrays.asList("execute", "executeQuery", "executeUpdate");
 
-    private final QueryCounterRequestContextScopeHolder queryCounterRequestContextScopeHolder;
+    private final QueryCountPerRequest queryCountPerRequest;
 
-    public PreparedStatementQueryMonitor(QueryCounterRequestContextScopeHolder queryCounterRequestContextScopeHolder) {
-        this.queryCounterRequestContextScopeHolder = queryCounterRequestContextScopeHolder;
+    public PreparedStatementQueryMonitor(QueryCountPerRequest queryCountPerRequest) {
+        this.queryCountPerRequest = queryCountPerRequest;
     }
 
     @Override
@@ -24,7 +24,7 @@ public Object invoke(MethodInvocation invocation) throws Throwable {
             Object result = invocation.proceed();
             long endTime = System.currentTimeMillis();
 
-            queryCounterRequestContextScopeHolder.get().incrementQueryCount(endTime - startTime);
+            queryCountPerRequest.incrementQueryCount(endTime - startTime);
 
             return result;
         }

src/main/java/hyeon9mak/multidatasourcequerycounter/QueryCountLogger.java

@@ -7,6 +7,8 @@
 @Component
 public class QueryCountLogger {
 
+    private static final String LOG_MESSAGE_FORMAT = "{} - totalQueryCount: {}, totalSpendTime: {}ms";
+
     private final Logger logger = LoggerFactory.getLogger(getClass());
     private final QueryCounterLoggingLevelProperties properties;
 
@@ -16,11 +18,11 @@ public QueryCountLogger(QueryCounterLoggingLevelProperties properties) {
 
     public void logQueryCount(QueryCountPerRequest queryCountPerRequest) {
         if (properties.getError().isEnable() && queryCountPerRequest.getTotalQueryCount() >= properties.getError().getCount()) {
-            logger.error("{}, totalQueryCount: {}, totalSpendTime: {}ms", queryCountPerRequest.getApiUrl(), queryCountPerRequest.getTotalQueryCount(), queryCountPerRequest.getTotalQueryMilliSeconds());
+            logger.error(LOG_MESSAGE_FORMAT, queryCountPerRequest.getApiUrl(), queryCountPerRequest.getTotalQueryCount(), queryCountPerRequest.getTotalQueryMilliSeconds());
         } else if (properties.getWarn().isEnable() && queryCountPerRequest.getTotalQueryCount() >= properties.getWarn().getCount()) {
-            logger.warn("{}, totalQueryCount: {}, totalSpendTime: {}ms", queryCountPerRequest.getApiUrl(), queryCountPerRequest.getTotalQueryCount(), queryCountPerRequest.getTotalQueryMilliSeconds());
+            logger.warn(LOG_MESSAGE_FORMAT, queryCountPerRequest.getApiUrl(), queryCountPerRequest.getTotalQueryCount(), queryCountPerRequest.getTotalQueryMilliSeconds());
         } else if (properties.getInfo().isEnable() && queryCountPerRequest.getTotalQueryCount() >= properties.getInfo().getCount()) {
-            logger.info("{}, totalQueryCount: {}, totalSpendTime: {}ms", queryCountPerRequest.getApiUrl(), queryCountPerRequest.getTotalQueryCount(), queryCountPerRequest.getTotalQueryMilliSeconds());
+            logger.info(LOG_MESSAGE_FORMAT, queryCountPerRequest.getApiUrl(), queryCountPerRequest.getTotalQueryCount(), queryCountPerRequest.getTotalQueryMilliSeconds());
         }
     }
 }