Drop AS token when creating JPQL count query using DISTINCT

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-jpa-parent</artifactId>
-	<version>3.3.0-SNAPSHOT</version>
+	<version>3.3.x-3269-SNAPSHOT</version>
 	<packaging>pom</packaging>
 
 	<name>Spring Data JPA Parent</name>

spring-data-envers/pom.xml

@@ -5,12 +5,12 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-envers</artifactId>
-	<version>3.3.0-SNAPSHOT</version>
+	<version>3.3.x-3269-SNAPSHOT</version>
 
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-jpa-parent</artifactId>
-		<version>3.3.0-SNAPSHOT</version>
+		<version>3.3.x-3269-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-jpa-distribution/pom.xml

@@ -14,7 +14,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-jpa-parent</artifactId>
-		<version>3.3.0-SNAPSHOT</version>
+		<version>3.3.x-3269-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-jpa/pom.xml

@@ -6,7 +6,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-jpa</artifactId>
-	<version>3.3.0-SNAPSHOT</version>
+	<version>3.3.x-3269-SNAPSHOT</version>
 
 	<name>Spring Data JPA</name>
 	<description>Spring Data module for JPA repositories.</description>
@@ -15,7 +15,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-jpa-parent</artifactId>
-		<version>3.3.0-SNAPSHOT</version>
+		<version>3.3.x-3269-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/EqlQueryTransformer.java

@@ -161,12 +161,14 @@ public List<JpaQueryParsingToken> visitSelect_clause(EqlParser.Select_clauseCont
 
 				if (ctx.DISTINCT() != null) {
 
-					if (selectItemTokens.stream().anyMatch(jpqlToken -> jpqlToken.getToken().contains("new"))) {
+					List<JpaQueryParsingToken> countSelection = QueryTransformers.filterCountSelection(selectItemTokens);
+
+					if (countSelection.stream().anyMatch(jpqlToken -> jpqlToken.getToken().contains("new"))) {
 						// constructor
 						tokens.add(new JpaQueryParsingToken(() -> primaryFromAlias));
 					} else {
 						// keep all the select items to distinct against
-						tokens.addAll(selectItemTokens);
+						tokens.addAll(countSelection);
 					}
 				} else {
 					tokens.add(new JpaQueryParsingToken(() -> primaryFromAlias));
@@ -240,4 +242,5 @@ public List<JpaQueryParsingToken> visitConstructor_expression(EqlParser.Construc
 	private static <T> ArrayList<T> newArrayList() {
 		return new ArrayList<>();
 	}
+
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/HqlQueryTransformer.java

@@ -30,6 +30,7 @@
  * An ANTLR {@link org.antlr.v4.runtime.tree.ParseTreeVisitor} that transforms a parsed HQL query.
  *
  * @author Greg Turnquist
+ * @author Christoph Strobl
  * @since 3.1
  */
 class HqlQueryTransformer extends HqlQueryRenderer {
@@ -357,12 +358,14 @@ public List<JpaQueryParsingToken> visitSelectClause(HqlParser.SelectClauseContex
 
 				if (ctx.DISTINCT() != null) {
 
-					if (selectionListTokens.stream().anyMatch(hqlToken -> hqlToken.getToken().contains("new"))) {
+					List<JpaQueryParsingToken> countSelection = QueryTransformers.filterCountSelection(selectionListTokens);
+
+					if (countSelection.stream().anyMatch(hqlToken -> hqlToken.getToken().contains("new"))) {
 						// constructor
 						tokens.add(new JpaQueryParsingToken(() -> primaryFromAlias));
 					} else {
 						// keep all the select items to distinct against
-						tokens.addAll(selectionListTokens);
+						tokens.addAll(countSelection);
 					}
 				} else {
 					tokens.add(new JpaQueryParsingToken(() -> primaryFromAlias));
@@ -394,4 +397,5 @@ public List<JpaQueryParsingToken> visitInstantiation(HqlParser.InstantiationCont
 	static <T> ArrayList<T> newArrayList() {
 		return new ArrayList<>();
 	}
+
 }

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpaQueryParsingToken.java

@@ -131,6 +131,15 @@ boolean getSpace() {
 		return this.space;
 	}
 
+	boolean isA(JpaQueryParsingToken token) {
+		return token.getToken().equalsIgnoreCase(this.getToken());
+	}
+
+	@Override
+	public String toString() {
+		return getToken();
+	}
+
 	/**
 	 * Switch the last {@link JpaQueryParsingToken}'s spacing to {@literal true}.
 	 */

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/JpqlQueryTransformer.java

@@ -161,12 +161,14 @@ public List<JpaQueryParsingToken> visitSelect_clause(JpqlParser.Select_clauseCon
 
 				if (ctx.DISTINCT() != null) {
 
-					if (selectItemTokens.stream().anyMatch(jpqlToken -> jpqlToken.getToken().contains("new"))) {
+					List<JpaQueryParsingToken> countSelection = QueryTransformers.filterCountSelection(selectItemTokens);
+
+					if (countSelection.stream().anyMatch(jpqlToken -> jpqlToken.getToken().contains("new"))) {
 						// constructor
 						tokens.add(new JpaQueryParsingToken(() -> primaryFromAlias));
 					} else {
 						// keep all the select items to distinct against
-						tokens.addAll(selectItemTokens);
+						tokens.addAll(countSelection);
 					}
 				} else {
 					tokens.add(new JpaQueryParsingToken(() -> primaryFromAlias));

spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/query/QueryTransformers.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2024 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
+ *
+ *      https://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.query;
+
+import static org.springframework.data.jpa.repository.query.JpaQueryParsingToken.*;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Utility class encapsulating common query transformations.
+ *
+ * @author Mark Paluch
+ */
+class QueryTransformers {
+
+	static List<JpaQueryParsingToken> filterCountSelection(List<JpaQueryParsingToken> selection) {
+
+		List<JpaQueryParsingToken> target = new ArrayList<>(selection.size());
+		boolean skipNext = false;
+
+		for (JpaQueryParsingToken token : selection) {
+
+			if (skipNext) {
+				skipNext = false;
+				continue;
+			}
+
+			if (token.isA(TOKEN_AS)) {
+				skipNext = true;
+				continue;
+			}
+			target.add(token);
+		}
+
+		return target;
+	}
+
+}

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/EqlQueryTransformerTests.java

@@ -615,6 +615,20 @@ void countProjectionDistinctQueryIncludesNewLineAfterEntityAndBeforeWhere() {
 				"SELECT count(DISTINCT entity1) FROM Entity1 entity1 LEFT JOIN entity1.entity2 entity2 ON entity1.key = entity2.key where entity1.id = 1799");
 	}
 
+	@Test // GH-3269
+	void createsCountQueryUsingAliasCorrectly() {
+
+		assertCountQuery("select distinct 1 as x from Employee e", "select count(distinct 1) from Employee e");
+		assertCountQuery("SELECT DISTINCT abc AS x FROM T t", "SELECT count(DISTINCT abc) FROM T t");
+		assertCountQuery("select distinct a as x, b as y from Employee e", "select count(distinct a , b) from Employee e");
+		assertCountQuery("select distinct sum(amount) as x from Employee e GROUP BY n",
+				"select count(distinct sum(amount)) from Employee e GROUP BY n");
+		assertCountQuery("select distinct a, b, sum(amount) as c, d from Employee e GROUP BY n",
+				"select count(distinct a, b, sum(amount) , d) from Employee e GROUP BY n");
+		assertCountQuery("select distinct a, count(b) as c from Employee e GROUP BY n",
+				"select count(distinct a, count(b)) from Employee e GROUP BY n");
+	}
+
 	@Test // GH-2393
 	void createCountQueryStartsWithWhitespace() {
 

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/HqlQueryTransformerTests.java

@@ -34,6 +34,7 @@
  * Verify that HQL queries are properly transformed through the {@link JpaQueryEnhancer} and the {@link HqlQueryParser}.
  *
  * @author Greg Turnquist
+ * @author Christoph Strobl
  */
 class HqlQueryTransformerTests {
 
@@ -1031,6 +1032,17 @@ void aliasesShouldNotOverlapWithSortProperties() {
 						"SELECT t3 FROM Test3 t3 JOIN t3.test2 x WHERE x.id = :test2Id order by t3.testDuplicateColumnName desc");
 	}
 
+	@Test // GH-3269
+	void createsCountQueryUsingAliasCorrectly() {
+
+		assertCountQuery("select distinct 1 as x from Employee","select count(distinct 1) from Employee AS __");
+		assertCountQuery("SELECT DISTINCT abc AS x FROM T","SELECT count(DISTINCT abc) FROM T AS __");
+		assertCountQuery("select distinct a as x, b as y from Employee","select count(distinct a , b) from Employee AS __");
+		assertCountQuery("select distinct sum(amount) as x from Employee GROUP BY n","select count(distinct sum(amount)) from Employee AS __ GROUP BY n");
+		assertCountQuery("select distinct a, b, sum(amount) as c, d from Employee GROUP BY n","select count(distinct a, b, sum(amount) , d) from Employee AS __ GROUP BY n");
+		assertCountQuery("select distinct a, count(b) as c from Employee GROUP BY n","select count(distinct a, count(b)) from Employee AS __ GROUP BY n");
+	}
+
 	private void assertCountQuery(String originalQuery, String countQuery) {
 		assertThat(createCountQueryFor(originalQuery)).isEqualTo(countQuery);
 	}

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/JpqlQueryTransformerTests.java

@@ -679,6 +679,20 @@ void countQueryUsesCorrectVariable() {
 						.isEqualTo("SELECT count(us) FROM users_statuses us WHERE (user_created_at BETWEEN :fromDate AND :toDate)");
 	}
 
+	@Test // GH-3269
+	void createsCountQueryUsingAliasCorrectly() {
+
+		assertCountQuery("select distinct 1 as x from Employee e", "select count(distinct 1) from Employee e");
+		assertCountQuery("SELECT DISTINCT abc AS x FROM T t", "SELECT count(DISTINCT abc) FROM T t");
+		assertCountQuery("select distinct a as x, b as y from Employee e", "select count(distinct a , b) from Employee e");
+		assertCountQuery("select distinct sum(amount) as x from Employee e GROUP BY n",
+				"select count(distinct sum(amount)) from Employee e GROUP BY n");
+		assertCountQuery("select distinct a, b, sum(amount) as c, d from Employee e GROUP BY n",
+				"select count(distinct a, b, sum(amount) , d) from Employee e GROUP BY n");
+		assertCountQuery("select distinct a, count(b) as c from Employee e GROUP BY n",
+				"select count(distinct a, count(b)) from Employee e GROUP BY n");
+	}
+
 	@Test // GH-2496, GH-2522, GH-2537, GH-2045
 	void orderByShouldWorkWithSubSelectStatements() {