Make sure to clear the ThreadLocal variables

Author: reda-alaoui

core/src/main/java/com/cosium/spring/data/jpa/entity/graph/repository/support/CountQueryDetector.java

@@ -41,7 +41,7 @@ public Object invoke(MethodInvocation invocation) throws Throwable {
     try {
       return invocation.proceed();
     } finally {
-      IS_COUNT_QUERY.set(false);
+      IS_COUNT_QUERY.remove();
     }
   }
 }

core/src/main/java/com/cosium/spring/data/jpa/entity/graph/repository/support/RepositoryMethodEntityGraphExtractor.java

@@ -84,7 +84,11 @@ public Object invoke(MethodInvocation invocation) throws Throwable {
       try {
         return doInvoke(invocation);
       } finally {
-        CURRENT_REPOSITORY.set(oldRepo);
+        if (oldRepo == null) {
+          CURRENT_REPOSITORY.remove();
+        } else {
+          CURRENT_REPOSITORY.set(oldRepo);
+        }
       }
     }
 
@@ -142,7 +146,11 @@ private Object doInvoke(MethodInvocation invocation) throws Throwable {
         return invocation.proceed();
       } finally {
         if (newEntityGraphCandidatePreValidated) {
-          currentEntityGraph.set(genuineCandidate);
+          if (genuineCandidate == null) {
+            currentEntityGraph.remove();
+          } else {
+            currentEntityGraph.set(genuineCandidate);
+          }
         }
       }
     }