Feature/filter for queries

CHANGELOG.md

@@ -54,3 +54,4 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [1.0.1]: https://github.com/yannbriancon/spring-hibernate-query-utils/compare/v1.0.0...v1.0.1
 [1.0.0]: https://github.com/yannbriancon/spring-hibernate-query-utils/compare/v0.1.0...v1.0.0
 [0.1.0]: https://github.com/yannbriancon/spring-hibernate-query-utils/tree/v0.1.0
+- Accept filter criteria for query counting

src/main/java/com/yannbriancon/interceptor/HibernateQueryInterceptor.java

@@ -11,36 +11,34 @@
 
 import javax.persistence.EntityManager;
 import java.io.Serializable;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
+import java.util.*;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 
 @Component
 @EnableConfigurationProperties(NPlusOneQueriesDetectionProperties.class)
 public class HibernateQueryInterceptor extends EmptyInterceptor {
 
-    private transient ThreadLocal<Long> threadQueryCount = new ThreadLocal<>();
+    private final transient ThreadLocal<Long> threadQueryCount = new ThreadLocal<>();
 
-    private transient ThreadLocal<Set<String>> threadPreviouslyLoadedEntities =
+    private final transient ThreadLocal<Set<String>> threadPreviouslyLoadedEntities =
             ThreadLocal.withInitial(new EmptySetSupplier<>());
 
-    private transient ThreadLocal<Map<String, SelectQueriesInfo>> threadSelectQueriesInfoPerProxyMethod =
+    private final transient ThreadLocal<Map<String, SelectQueriesInfo>> threadSelectQueriesInfoPerProxyMethod =
             ThreadLocal.withInitial(new EmptyMapSupplier<>());
 
     private static final Logger LOGGER = LoggerFactory.getLogger(HibernateQueryInterceptor.class);
 
-    private final NPlusOneQueriesDetectionProperties NPlusOneQueriesDetectionProperties;
+    private final NPlusOneQueriesDetectionProperties nPlusOneQueriesDetectionProperties;
 
     private static final String HIBERNATE_PROXY_PREFIX = "org.hibernate.proxy";
     private static final String PROXY_METHOD_PREFIX = "com.sun.proxy";
+    private transient ThreadLocal<Predicate<String>> queryFilter;
 
     public HibernateQueryInterceptor(
-            NPlusOneQueriesDetectionProperties NPlusOneQueriesDetectionProperties
+            NPlusOneQueriesDetectionProperties nPlusOneQueriesDetectionProperties
     ) {
-        this.NPlusOneQueriesDetectionProperties = NPlusOneQueriesDetectionProperties;
+        this.nPlusOneQueriesDetectionProperties = nPlusOneQueriesDetectionProperties;
     }
 
     /**
@@ -67,6 +65,18 @@ public void clearNPlusOneQuerySession(EntityManager entityManager) {
      * Start or reset the query count to 0 for the considered thread
      */
     public void startQueryCount() {
+        startQueryCount(null);
+    }
+
+    /**
+     * Start or reset the query count to 0 for the considered thread.
+     * Counts only if given filter matches.
+     *
+     * @param filter string filter predicate
+     */
+    public void startQueryCount(Predicate<String> filter) {
+        this.queryFilter = new ThreadLocal<>();
+        this.queryFilter.set(filter);
         threadQueryCount.set(0L);
     }
 
@@ -79,18 +89,18 @@ public Long getQueryCount() {
 
     /**
      * Detect the N+1 queries by keeping the history of sql queries generated per proxy method.
-     * Increment the query count for the considered thread for each new statement if the count has been initialized.
+     * Increment the query count for the considered thread for each new statement if the count has been initialized and the filter matches.
      *
      * @param sql Query to be executed
      * @return Query to be executed
      */
     @Override
     public String onPrepareStatement(String sql) {
-        if (NPlusOneQueriesDetectionProperties.isEnabled()) {
+        if (nPlusOneQueriesDetectionProperties.isEnabled()) {
             updateSelectQueriesInfoPerProxyMethod(sql);
         }
         Long count = threadQueryCount.get();
-        if (count != null) {
+        if (count != null && (Objects.isNull(queryFilter.get()) || queryFilter.get().test(sql))) {
             threadQueryCount.set(count + 1);
         }
         return super.onPrepareStatement(sql);
@@ -115,7 +125,7 @@ public void afterTransactionCompletion(Transaction tx) {
      */
     @Override
     public Object getEntity(String entityName, Serializable id) {
-        if (NPlusOneQueriesDetectionProperties.isEnabled()) {
+        if (nPlusOneQueriesDetectionProperties.isEnabled()) {
             detectNPlusOneQueriesFromMissingEagerFetchingOnAQuery(entityName, id);
             detectNPlusOneQueriesFromClassFieldEagerFetching(entityName);
         }
@@ -235,23 +245,23 @@ private void detectNPlusOneQueriesFromClassFieldEagerFetching(String entityName)
         selectQueriesInfoPerProxyMethod.put(proxyMethodName, selectQueriesInfo.resetSelectQueriesCount());
         threadSelectQueriesInfoPerProxyMethod.set(selectQueriesInfoPerProxyMethod);
 
-        String errorMessage = "N+1 queries detected with eager fetching on the entity " + entityName;
+        StringBuilder errorMessage = new StringBuilder();
+        errorMessage.append("N+1 queries detected with eager fetching on the entity ").append(entityName);
 
         // Find origin of the N+1 queries in client package
-        // by getting oldest occurrence of proxy method in stack elements
+        // by getting the oldest occurrence of proxy method in stack elements
         StackTraceElement[] stackTraceElements = Thread.currentThread().getStackTrace();
 
         for (int i = stackTraceElements.length - 1; i >= 1; i--) {
             if (stackTraceElements[i - 1].getClassName().indexOf(PROXY_METHOD_PREFIX) == 0) {
-                errorMessage += "\n    at " + stackTraceElements[i].toString();
+                errorMessage.append("\n    at ").append(stackTraceElements[i].toString());
                 break;
             }
         }
 
-        errorMessage += "\n    Hint: Missing Lazy fetching configuration on a field of type " + entityName + " of " +
-                "one of the entities fetched in the query\n";
+        errorMessage.append("\n    Hint: Missing Lazy fetching configuration on a field of type ").append(entityName).append(" of ").append("one of the entities fetched in the query\n");
 
-        logDetectedNPlusOneQueries(errorMessage);
+        logDetectedNPlusOneQueries(errorMessage.toString());
     }
 
     /**
@@ -279,7 +289,7 @@ private Optional<String> getProxyMethodName() {
      * @param errorMessage Error message for the N+1 queries detected
      */
     private void logDetectedNPlusOneQueries(String errorMessage) {
-        switch (NPlusOneQueriesDetectionProperties.getErrorLevel()) {
+        switch (nPlusOneQueriesDetectionProperties.getErrorLevel()) {
             case INFO:
                 LOGGER.info(errorMessage);
                 break;

src/test/java/com/yannbriancon/interceptor/QueryCountTest.java

@@ -4,13 +4,20 @@
 import com.yannbriancon.utils.repository.MessageRepository;
 import com.yannbriancon.utils.repository.UserRepository;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.runner.RunWith;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.test.context.junit4.SpringRunner;
 import org.springframework.transaction.annotation.Transactional;
 
+import java.util.function.Predicate;
+import java.util.stream.Stream;
+
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.params.provider.Arguments.arguments;
 
 @RunWith(SpringRunner.class)
 @SpringBootTest
@@ -45,4 +52,34 @@ void queryCount_isOkWhenSaveQueryIsExecutedBeforeStartingTheCount() {
 
         assertThat(hibernateQueryInterceptor.getQueryCount()).isEqualTo(1);
     }
+
+    @Test
+    void queryCount_WithoutFilter() {
+        hibernateQueryInterceptor.startQueryCount();
+        userRepository.saveAndFlush(new User());
+
+        messageRepository.findAll();
+
+        assertThat(hibernateQueryInterceptor.getQueryCount()).isEqualTo(2);
+    }
+
+    @ParameterizedTest
+    @MethodSource("prepareQueryCountWithStartsWithFilterData")
+    void queryCountWithStartsWithFilter(String filter, int expectedQueryCount) {
+        final Predicate<String> predicate = f -> f.startsWith(filter);
+
+        hibernateQueryInterceptor.startQueryCount(predicate);
+
+        userRepository.saveAndFlush(new User());
+        messageRepository.findAll();
+
+        assertThat(hibernateQueryInterceptor.getQueryCount()).isEqualTo(expectedQueryCount);
+    }
+
+    static Stream<Arguments> prepareQueryCountWithStartsWithFilterData() {
+        return Stream.of(arguments("select", 1),
+                arguments("insert", 1),
+                arguments("update", 0),
+                arguments("delete", 0));
+    }
 }