From 0aa6d7d6e79554e23cdfe65c18565735cac78a99 Mon Sep 17 00:00:00 2001
From: "Greg L. Turnquist" <gturnquist@vmware.com>
Date: Tue, 10 May 2022 10:34:10 -0500
Subject: [PATCH] Apply query hints to count queries for page-based
 Specification.

Query hints are applied in many places, but not when doing a findAll(Specification, Pageable).

Closes #2054.
Original pull request #2528
---
 .../support/SimpleJpaRepository.java          | 24 ++++++++++++++++++-
 .../support/SimpleJpaRepositoryUnitTests.java | 17 ++++++++++---
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/support/SimpleJpaRepository.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/support/SimpleJpaRepository.java
index d47fd9883..9f4552d96 100644
--- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/support/SimpleJpaRepository.java
+++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/support/SimpleJpaRepository.java
@@ -309,6 +309,13 @@ public class SimpleJpaRepository<T, ID> implements JpaRepositoryImplementation<T
 		return metadata == null ? NoHints.INSTANCE : DefaultQueryHints.of(entityInformation, metadata);
 	}
 
+	/**
+	 * Returns {@link QueryHints} with the query hints on the current {@link CrudMethodMetadata} for count queries.
+	 */
+	protected QueryHints getQueryHintsForCount() {
+		return metadata == null ? NoHints.INSTANCE : DefaultQueryHints.of(entityInformation, metadata).forCounts();
+	}
+
 	@Deprecated
 	@Override
 	public T getOne(ID id) {
@@ -750,7 +757,7 @@ public class SimpleJpaRepository<T, ID> implements JpaRepositoryImplementation<T
 		// Remove all Orders the Specifications might have applied
 		query.orderBy(Collections.emptyList());
 
-		return em.createQuery(query);
+		return applyRepositoryMethodMetadataForCount(em.createQuery(query));
 	}
 
 	/**
@@ -800,6 +807,21 @@ public class SimpleJpaRepository<T, ID> implements JpaRepositoryImplementation<T
 		getQueryHints().withFetchGraphs(em).forEach(query::setHint);
 	}
 
+	private <S> TypedQuery<S> applyRepositoryMethodMetadataForCount(TypedQuery<S> query) {
+
+		if (metadata == null) {
+			return query;
+		}
+
+		applyQueryHintsForCount(query);
+
+		return query;
+	}
+
+	private void applyQueryHintsForCount(Query query) {
+		getQueryHintsForCount().forEach(query::setHint);
+	}
+
 	/**
 	 * Executes a count query and transparently sums up all values returned.
 	 *
diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/support/SimpleJpaRepositoryUnitTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/support/SimpleJpaRepositoryUnitTests.java
index cc69535d3..ebbc54904 100644
--- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/support/SimpleJpaRepositoryUnitTests.java
+++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/support/SimpleJpaRepositoryUnitTests.java
@@ -18,9 +18,7 @@ package org.springframework.data.jpa.repository.support;
 import static java.util.Collections.*;
 import static org.assertj.core.api.Assertions.*;
 import static org.mockito.Mockito.*;
-
-import java.util.Arrays;
-import java.util.Optional;
+import static org.springframework.data.jpa.domain.Specification.*;
 
 import jakarta.persistence.EntityGraph;
 import jakarta.persistence.EntityManager;
@@ -28,6 +26,9 @@ import jakarta.persistence.TypedQuery;
 import jakarta.persistence.criteria.CriteriaBuilder;
 import jakarta.persistence.criteria.CriteriaQuery;
 
+import java.util.Arrays;
+import java.util.Optional;
+
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -192,4 +193,14 @@ class SimpleJpaRepositoryUnitTests {
 		verify(em, never()).remove(newUser);
 		verify(em, never()).merge(newUser);
 	}
+
+	@Test // GH-2054
+	void applyQueryHintsToCountQueriesForSpecificationPageables() {
+
+		when(query.getResultList()).thenReturn(Arrays.asList(new User(), new User()));
+
+		repo.findAll(where(null), PageRequest.of(2, 1));
+
+		verify(metadata).getQueryHintsForCount();
+	}
 }