From 9f7c591963d6570138cf8cd329406fe5fc6a78f6 Mon Sep 17 00:00:00 2001
From: Mark Paluch <mark.paluch@broadcom.com>
Date: Tue, 10 Dec 2024 15:17:37 +0100
Subject: [PATCH] Use primary table alias for SQL count query derivation with
 DISTINCT queries.

We now use render COUNT(DISTINCT a.*) where 'a' is the primary table alias instead of COUNT(DISTINCT *).

Closes #3707
---
 .../query/JSqlParserQueryEnhancer.java         | 18 ++++++++++++------
 .../JSqlParserQueryEnhancerUnitTests.java      | 16 ++++++++++++++++
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JSqlParserQueryEnhancer.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JSqlParserQueryEnhancer.java
index a76e04d5e7..1fbc87d17a 100644
--- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JSqlParserQueryEnhancer.java
+++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JSqlParserQueryEnhancer.java
@@ -58,6 +58,7 @@
  * @author Geoffrey Deremetz
  * @author Yanming Zhou
  * @author Christoph Strobl
+ * @author Mark Paluch
  * @since 2.7.0
  */
 public class JSqlParserQueryEnhancer implements QueryEnhancer {
@@ -119,7 +120,7 @@ public String applySorting(Sort sort, @Nullable String alias) {
 
 		Select selectStatement = parseSelectStatement(queryString);
 
-		if (selectStatement  instanceof SetOperationList setOperationList) {
+		if (selectStatement instanceof SetOperationList setOperationList) {
 			return applySortingToSetOperationList(setOperationList, sort);
 		}
 
@@ -217,7 +218,7 @@ private Set<String> getJoinAliases(String query) {
 		}
 
 		Select selectStatement = (Select) statement;
-		if (selectStatement  instanceof PlainSelect selectBody) {
+		if (selectStatement instanceof PlainSelect selectBody) {
 			return getJoinAliases(selectBody);
 		}
 
@@ -315,7 +316,7 @@ private String detectAlias(String query) {
 			 * ValuesStatement has no alias
 			 * SetOperation can have multiple alias for each operation item
 			 */
-			if (!(selectStatement  instanceof PlainSelect selectBody)) {
+			if (!(selectStatement instanceof PlainSelect selectBody)) {
 				return null;
 			}
 
@@ -370,7 +371,7 @@ public String createCountQueryFor(@Nullable String countProjection) {
 		/*
 		  We only support count queries for {@link PlainSelect}.
 		 */
-		if (!(selectStatement  instanceof PlainSelect selectBody)) {
+		if (!(selectStatement instanceof PlainSelect selectBody)) {
 			return this.query.getQueryString();
 		}
 
@@ -413,7 +414,7 @@ public String getProjection() {
 
 		Select selectBody = selectStatement;
 
-		if (selectStatement  instanceof SetOperationList setOperationList) {
+		if (selectStatement instanceof SetOperationList setOperationList) {
 
 			// using the first one since for setoperations the projection has to be the same
 			selectBody = setOperationList.getSelects().get(0);
@@ -491,7 +492,12 @@ private String countPropertyNameForSelection(List<SelectItem<?>> selectItems, bo
 			return column.getFullyQualifiedName();
 		}
 
-		return query.isNativeQuery() ? (distinct ? "*" : "1") : tableAlias == null ? "*" : tableAlias;
+		// TODO: We should not handle JPQL here...
+		if (!query.isNativeQuery()) {
+			return tableAlias == null ? "*" : tableAlias;
+		}
+
+		return distinct ? ((tableAlias != null ? tableAlias + "." : "") + "*") : "1";
 	}
 
 	@Override
diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JSqlParserQueryEnhancerUnitTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JSqlParserQueryEnhancerUnitTests.java
index a1f3015b59..1269fe23f0 100644
--- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JSqlParserQueryEnhancerUnitTests.java
+++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JSqlParserQueryEnhancerUnitTests.java
@@ -51,6 +51,22 @@ void shouldApplySorting() {
 		assertThat(sql).isEqualTo("SELECT e FROM Employee e ORDER BY e.foo ASC, e.bar ASC");
 	}
 
+	@Test // GH-3707
+	void countQueriesShouldConsiderPrimaryTableAlias() {
+
+		QueryEnhancer enhancer = createQueryEnhancer(DeclaredQuery.of("""
+				SELECT DISTINCT a.*, b.b1
+				FROM TableA a
+				  JOIN TableB b ON a.b = b.b
+				  LEFT JOIN TableC c ON b.c = c.c
+				ORDER BY b.b1, a.a1, a.a2
+				""", true));
+
+		String sql = enhancer.createCountQueryFor();
+
+		assertThat(sql).startsWith("SELECT count(DISTINCT a.*) FROM TableA a");
+	}
+
 	@Override
 	@ParameterizedTest // GH-2773
 	@MethodSource("jpqlCountQueries")