From 9d4bde3e7d4c30908ce98750eb12c3afa99fe0e8 Mon Sep 17 00:00:00 2001
From: Mark Paluch <mark.paluch@broadcom.com>
Date: Thu, 1 Aug 2024 09:50:01 +0200
Subject: [PATCH] Suppress selection item aliasing for the actual count query.

We now no longer apply count selection filtering but rather skip select field aliasing when rendering a count query to drop the field alias within a count query.

Previously, we removed field aliasing by filtering the token stream which also removed the AS keyword from cast operators.

Closes: #3536
Original Pull Request: #3553
---
 .../query/HqlCountQueryTransformer.java       | 25 +++++++++++++------
 .../query/HqlQueryTransformerTests.java       | 17 +++++++++++++
 2 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlCountQueryTransformer.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlCountQueryTransformer.java
index 537da47e71..4e083c757d 100644
--- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlCountQueryTransformer.java
+++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlCountQueryTransformer.java
@@ -15,12 +15,7 @@
  */
 package org.springframework.data.jpa.repository.query;
 
-import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_AS;
-import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_CLOSE_PAREN;
-import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_COUNT_FUNC;
-import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_DOUBLE_UNDERSCORE;
-import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_OPEN_PAREN;
-import static org.springframework.data.jpa.repository.query.QueryTokens.TOKEN_SELECT_COUNT;
+import static org.springframework.data.jpa.repository.query.QueryTokens.*;
 
 import org.springframework.data.jpa.repository.query.HqlParser.SelectClauseContext;
 import org.springframework.data.jpa.repository.query.QueryRenderer.QueryRendererBuilder;
@@ -208,6 +203,22 @@ public QueryTokenStream visitSelectClause(HqlParser.SelectClauseContext ctx) {
 		return builder;
 	}
 
+	@Override
+	public QueryTokenStream visitSelection(HqlParser.SelectionContext ctx) {
+
+		if (isSubquery(ctx)) {
+			return super.visitSelection(ctx);
+		}
+
+		QueryRendererBuilder builder = QueryRenderer.builder();
+
+		builder.append(visit(ctx.selectExpression()));
+
+		// do not append variables to skip AS field aliasing
+
+		return builder;
+	}
+
 	@Override
 	public QueryRendererBuilder visitQueryOrder(HqlParser.QueryOrderContext ctx) {
 
@@ -247,7 +258,7 @@ private QueryRendererBuilder getDistinctCountSelection(QueryTokenStream selectio
 			nested.append(QueryTokens.token(primaryFromAlias));
 		} else {
 			// keep all the select items to distinct against
-			nested.append(countSelection);
+			nested.append(selectionListbuilder);
 		}
 		return nested;
 	}
diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/HqlQueryTransformerTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/HqlQueryTransformerTests.java
index e19036dace..a9b692cd76 100644
--- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/HqlQueryTransformerTests.java
+++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/HqlQueryTransformerTests.java
@@ -90,6 +90,23 @@ void applyCountToSimpleQuery() {
 		assertThat(results).isEqualTo("select count(e) FROM Employee e where e.name = :name");
 	}
 
+	@Test // GH-3536
+	void shouldCreateCountQueryForDistinctCount() {
+
+		// given
+		var original = """
+				select distinct cast(e.timestampField as date) as foo
+				from ExampleEntity e
+				order by cast(e.timestampField as date) desc
+				""";
+
+		// when
+		var results = createCountQueryFor(original);
+
+		// then
+		assertThat(results).isEqualTo("select count(distinct cast(e.timestampField as date)) from ExampleEntity e");
+	}
+
 	@Test
 	void applyCountToMoreComplexQuery() {