implement sql query filter. fix #14

Author: samuelstein

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [2.0.0]
 ### Added
+- Accept filter criteria for query counting
 - Make autocomplete work for N+1 queries detection properties.
 - Add more precise detection to
     - Avoid false positive with multiple calls to same method

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

@@ -11,11 +11,8 @@
 
 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
@@ -36,6 +33,7 @@ public class HibernateQueryInterceptor extends EmptyInterceptor {
 
     private static final String HIBERNATE_PROXY_PREFIX = "org.hibernate.proxy";
     private static final String PROXY_METHOD_PREFIX = "com.sun.proxy";
+    private transient Predicate<String> filter;
 
     public HibernateQueryInterceptor(
             NPlusOneQueriesDetectionProperties NPlusOneQueriesDetectionProperties
@@ -70,16 +68,29 @@ public void startQueryCount() {
         threadQueryCount.set(0L);
     }
 
+    /**
+     * 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.filter = filter;
+        startQueryCount();
+    }
+
     /**
      * Get the query count for the considered thread
      */
     public Long getQueryCount() {
+        this.filter = null;
         return threadQueryCount.get();
     }
 
     /**
      * 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.
+     * If a query filter is set it will be preferred.
      *
      * @param sql Query to be executed
      * @return Query to be executed
@@ -91,7 +102,13 @@ public String onPrepareStatement(String sql) {
         }
         Long count = threadQueryCount.get();
         if (count != null) {
-            threadQueryCount.set(count + 1);
+            if (Objects.nonNull(filter)) {
+                if (filter.test(sql)) {
+                    threadQueryCount.set(count + 1);
+                }
+            } else {
+                threadQueryCount.set(count + 1);
+            }
         }
         return super.onPrepareStatement(sql);
     }

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,24 @@ void queryCount_isOkWhenSaveQueryIsExecutedBeforeStartingTheCount() {
 
         assertThat(hibernateQueryInterceptor.getQueryCount()).isEqualTo(1);
     }
+
+    @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));
+    }
 }