DATAJPA-1156 - Filter QueryHints for count queries in QuerydslJpaPredicateExecutor.

Original pull request: #279.

src/main/java/org/springframework/data/jpa/repository/support/CrudMethodMetadata.java

@@ -32,6 +32,7 @@ import org.springframework.lang.Nullable;
  * @author Thomas Darimont
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Jens Schauder
  */
 public interface CrudMethodMetadata {
 
@@ -50,6 +51,15 @@ public interface CrudMethodMetadata {
 	 */
 	Map<String, Object> getQueryHints();
 
+	/**
+	 * Returns all query hints to be applied to count queries executed for the CRUD method.
+	 *
+	 * The default implementation just delegates to {@link #getQueryHints()}.
+	 */
+	default Map<String, Object> getQueryHintsForCount() {
+		return getQueryHints();
+	}
+
 	/**
 	 * Returns the {@link EntityGraph} to be used.
 	 *

src/main/java/org/springframework/data/jpa/repository/support/CrudMethodMetadataPostProcessor.java

@@ -53,6 +53,7 @@ import org.springframework.util.ClassUtils;
  * @author Thomas Darimont
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Jens Schauder
  */
 class CrudMethodMetadataPostProcessor implements RepositoryProxyPostProcessor, BeanClassLoaderAware {
 
@@ -150,6 +151,7 @@ class CrudMethodMetadataPostProcessor implements RepositoryProxyPostProcessor, B
 
 		private final @Nullable LockModeType lockModeType;
 		private final Map<String, Object> queryHints;
+		private final Map<String, Object> getQueryHintsForCount;
 		private final Optional<EntityGraph> entityGraph;
 		private final Method method;
 
@@ -163,7 +165,8 @@ class CrudMethodMetadataPostProcessor implements RepositoryProxyPostProcessor, B
 			Assert.notNull(method, "Method must not be null!");
 
 			this.lockModeType = findLockModeType(method);
-			this.queryHints = findQueryHints(method);
+			this.queryHints = findQueryHints(method, false);
+			this.getQueryHintsForCount = findQueryHints(method, true);
 			this.entityGraph = findEntityGraph(method);
 			this.method = method;
 		}
@@ -179,12 +182,14 @@ class CrudMethodMetadataPostProcessor implements RepositoryProxyPostProcessor, B
 			return annotation == null ? null : (LockModeType) AnnotationUtils.getValue(annotation);
 		}
 
-		private static Map<String, Object> findQueryHints(Method method) {
+		private static Map<String, Object> findQueryHints(Method method, boolean forCount) {
 
 			Map<String, Object> queryHints = new HashMap<String, Object>();
 			QueryHints queryHintsAnnotation = AnnotatedElementUtils.findMergedAnnotation(method, QueryHints.class);
 
-			if (queryHintsAnnotation != null) {
+			if (queryHintsAnnotation != null
+					&& (!forCount || queryHintsAnnotation.forCounting())
+					) {
 
 				for (QueryHint hint : queryHintsAnnotation.value()) {
 					queryHints.put(hint.name(), hint.value());
@@ -219,6 +224,15 @@ class CrudMethodMetadataPostProcessor implements RepositoryProxyPostProcessor, B
 			return queryHints;
 		}
 
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.jpa.repository.support.CrudMethodMetadata#getQueryHintsForCount()
+		 */
+		@Override
+		public Map<String, Object> getQueryHintsForCount() {
+			return getQueryHintsForCount;
+		}
+
 		/*
 		 * (non-Javadoc)
 		 * @see org.springframework.data.jpa.repository.support.CrudMethodMetadata#getEntityGraph()

src/main/java/org/springframework/data/jpa/repository/support/DefaultQueryHints.java

@@ -20,6 +20,7 @@ import org.springframework.util.Assert;
  *
  * @author Christoph Strobl
  * @author Oliver Gierke
+ * @author Jens Schauder
  * @since 2.0
  */
 class DefaultQueryHints implements QueryHints {
@@ -27,6 +28,7 @@ class DefaultQueryHints implements QueryHints {
 	private final JpaEntityInformation<?, ?> information;
 	private final CrudMethodMetadata metadata;
 	private final Optional<EntityManager> entityManager;
+	private final boolean forCounts;
 
 	/**
 	 * Creates a new {@link DefaultQueryHints} instance for the given {@link JpaEntityInformation},
@@ -35,13 +37,15 @@ class DefaultQueryHints implements QueryHints {
 	 * @param information must not be {@literal null}.
 	 * @param metadata must not be {@literal null}.
 	 * @param entityManager must not be {@literal null}.
+	 * @param forCounts
 	 */
 	private DefaultQueryHints(JpaEntityInformation<?, ?> information, CrudMethodMetadata metadata,
-			Optional<EntityManager> entityManager) {
+							  Optional<EntityManager> entityManager, boolean forCounts) {
 
 		this.information = information;
 		this.metadata = metadata;
 		this.entityManager = entityManager;
+		this.forCounts = forCounts;
 	}
 
 	/**
@@ -57,7 +61,7 @@ class DefaultQueryHints implements QueryHints {
 		Assert.notNull(information, "JpaEntityInformation must not be null!");
 		Assert.notNull(metadata, "CrudMethodMetadata must not be null!");
 
-		return new DefaultQueryHints(information, metadata, Optional.empty());
+		return new DefaultQueryHints(information, metadata, Optional.empty(), false);
 	}
 
 	/*
@@ -66,7 +70,12 @@ class DefaultQueryHints implements QueryHints {
 	 */
 	@Override
 	public QueryHints withFetchGraphs(EntityManager em) {
-		return new DefaultQueryHints(this.information, this.metadata, Optional.of(em));
+		return new DefaultQueryHints(this.information, this.metadata, Optional.of(em), forCounts);
+	}
+
+	@Override
+	public QueryHints forCounts() {
+		return new DefaultQueryHints(this.information, this.metadata, entityManager, true);
 	}
 
 	/*
@@ -87,7 +96,12 @@ class DefaultQueryHints implements QueryHints {
 
 		Map<String, Object> hints = new HashMap<>();
 
-		hints.putAll(metadata.getQueryHints());
+		if (forCounts) {
+			hints.putAll(metadata.getQueryHintsForCount());
+		} else {
+			hints.putAll(metadata.getQueryHints());
+		}
+
 		hints.putAll(getFetchGraphs());
 
 		return hints;

src/main/java/org/springframework/data/jpa/repository/support/QueryHints.java

@@ -25,6 +25,14 @@ interface QueryHints extends Iterable<Entry<String, Object>> {
 	 */
 	QueryHints withFetchGraphs(EntityManager em);
 
+
+	/**
+	 * Creates and returns a new {@link QueryHints} instance that will contain only those hints applicable for count queries.
+	 *
+	 * @return new instance of {@link QueryHints}.
+	 */
+	QueryHints forCounts();
+
 	/**
 	 * Get the query hints as a {@link Map}.
 	 *
@@ -68,5 +76,14 @@ interface QueryHints extends Iterable<Entry<String, Object>> {
 		public QueryHints withFetchGraphs(EntityManager em) {
 			return this;
 		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.jpa.repository.support.QueryHints#forCounts(javax.persistence.EntityManager)
+		 */
+		@Override
+		public QueryHints forCounts() {
+			return this;
+		}
 	}
 }

src/main/java/org/springframework/data/jpa/repository/support/QuerydslJpaPredicateExecutor.java

@@ -196,7 +196,7 @@ public class QuerydslJpaPredicateExecutor<T> implements QuerydslPredicateExecuto
 	 * @return the Querydsl count {@link JPQLQuery}.
 	 */
 	protected JPQLQuery<?> createCountQuery(@Nullable Predicate... predicate) {
-		return doCreateQuery(getQueryHints(), predicate);
+		return doCreateQuery(getQueryHintsForCount(), predicate);
 	}
 
 	@Nullable
@@ -214,6 +214,16 @@ public class QuerydslJpaPredicateExecutor<T> implements QuerydslPredicateExecuto
 		return metadata == null ? QueryHints.NoHints.INSTANCE : DefaultQueryHints.of(entityInformation, metadata);
 	}
 
+	/**
+	 * Returns {@link QueryHints} with the query hints based on the current {@link CrudMethodMetadata} and potential
+	 * {@link EntityGraph} information and filtered for those hints that are to be applied to count queries.
+	 *
+	 * @return
+	 */
+	private QueryHints getQueryHintsForCount() {
+		return metadata == null ? QueryHints.NoHints.INSTANCE : DefaultQueryHints.of(entityInformation, metadata).forCounts();
+	}
+
 	private AbstractJPAQuery<?, ?> doCreateQuery(QueryHints hints, @Nullable Predicate... predicate) {
 
 		AbstractJPAQuery<?, ?> query = querydsl.createQuery(path);

src/test/java/org/springframework/data/jpa/repository/support/DefaultQueryHintsTest.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.jpa.repository.support;
+
+import static org.assertj.core.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+
+/**
+ * Unit tests for {@link DefaultQueryHints}.
+ *
+ * @author Jens Schauder
+ */
+public class DefaultQueryHintsTest {
+
+	JpaEntityInformation<?, ?> information = mock(JpaEntityInformation.class);
+	CrudMethodMetadata metadata = mock(CrudMethodMetadata.class);
+
+	@Test // DATAJPA-1156
+	public void mainHints() {
+
+		QueryHints hints = DefaultQueryHints.of(information, metadata);
+
+		assertThat(hints.asMap()) //
+				.extracting("name1", "name2", "n1", "n2") //
+				.containsExactly("value1", "value2", null, null);
+	}
+
+	@Test // DATAJPA-1156
+	public void countHints() {
+
+		QueryHints hints = DefaultQueryHints.of(information, metadata).forCounts();
+
+		assertThat(hints.asMap()) //
+				.extracting("name1", "name2", "n1", "n2") //
+				.containsExactly(null, null, "1", "2");
+	}
+
+
+	@Before
+	public void before() {
+
+		setupMainHints();
+		setUpCountHints();
+	}
+
+	private void setupMainHints() {
+		Map<String, Object> mainHintMap = new HashMap<>();
+		mainHintMap.put("name1", "value1");
+		mainHintMap.put("name2", "value2");
+		when(metadata.getQueryHints()).thenReturn(mainHintMap);
+	}
+
+	private void setUpCountHints() {
+		Map<String, Object> countHintMap = new HashMap<>();
+		countHintMap.put("n1", "1");
+		countHintMap.put("n2", "2");
+		when(metadata.getQueryHintsForCount()).thenReturn(countHintMap);
+	}
+
+}