From 0f6736fc1aa42debcdc687dfbeadb57fc6a41a35 Mon Sep 17 00:00:00 2001
From: Christoph Strobl <christoph.strobl@broadcom.com>
Date: Thu, 13 Jun 2024 11:01:16 +0200
Subject: [PATCH] Fix count query creation for HQL when using CTE.

Closes #3504
Original pull request #3508
---
 .../repository/query/HqlQueryTransformer.java |  2 +-
 .../query/HqlQueryTransformerTests.java       | 30 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlQueryTransformer.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlQueryTransformer.java
index f1e18f397..befaccc7b 100644
--- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlQueryTransformer.java
+++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlQueryTransformer.java
@@ -93,7 +93,7 @@ class HqlQueryTransformer extends HqlQueryRenderer {
 	 */
 	private static boolean isSubquery(ParserRuleContext ctx) {
 
-		if (ctx instanceof HqlParser.SubqueryContext) {
+		if (ctx instanceof HqlParser.SubqueryContext || ctx instanceof HqlParser.CteContext) {
 			return true;
 		} else if (ctx instanceof HqlParser.SelectStatementContext) {
 			return false;
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 c476f0d5b..0c4b3251f 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
@@ -500,6 +500,36 @@ class HqlQueryTransformerTests {
 				""");
 	}
 
+	@Test // GH-3504
+	void createCountWithCteShouldWork() {
+
+		String countQuery = createCountQueryFor("""
+				WITH maxId AS(select max(sr.snapshot.id) snapshotId from SnapshotReference sr
+					where sr.id.selectionId = ?1 and sr.enabled
+					group by sr.userId)
+				select sr from maxId m join SnapshotReference sr on sr.snapshot.id = m.snapshotId
+				""");
+
+		assertThat(countQuery).startsWith("WITH maxId AS(select max(sr.snapshot.id) snapshotId from SnapshotReference sr")
+				.endsWith("select count(m) from maxId m join SnapshotReference sr on sr.snapshot.id = m.snapshotId");
+	}
+
+	@Test // GH-3504
+	void createSortedQueryWithCteShouldWork() {
+
+		String sortedQuery = createQueryFor("""
+				WITH maxId AS(select max(sr.snapshot.id) snapshotId from SnapshotReference sr
+					where sr.id.selectionId = ?1 and sr.enabled
+					group by sr.userId)
+				select sr from maxId m join SnapshotReference sr on sr.snapshot.id = m.snapshotId
+				""", Sort.by("sr.snapshot"));
+
+		assertThat(sortedQuery).startsWith(
+				"WITH maxId AS(select max(sr.snapshot.id) snapshotId from SnapshotReference sr where sr.id.selectionId = ?1 and sr.enabled group by sr.userId )")
+				.endsWith(
+						"select sr from maxId m join SnapshotReference sr on sr.snapshot.id = m.snapshotId order by sr.snapshot asc");
+	}
+
 	@Test
 	void createCountQuerySupportsLineBreaksInSelectClause() {