Polishing.

Apply select list rewriting to EQL and JPQL parsers as well.

See: #3269
Original pull request: #3276

Author: mp911de

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

@@ -358,7 +358,7 @@ public List<JpaQueryParsingToken> visitSelectClause(HqlParser.SelectClauseContex
 
 				if (ctx.DISTINCT() != null) {
 
-					List<JpaQueryParsingToken> countSelection = getCountSelection(selectionListTokens);
+					List<JpaQueryParsingToken> countSelection = QueryTransformers.filterCountSelection(selectionListTokens);
 
 					if (countSelection.stream().anyMatch(hqlToken -> hqlToken.getToken().contains("new"))) {
 						// constructor
@@ -398,18 +398,4 @@ static <T> ArrayList<T> newArrayList() {
 		return new ArrayList<>();
 	}
 
-	private static List<JpaQueryParsingToken> getCountSelection(List<JpaQueryParsingToken> selectionListTokens) {
-
-		List<JpaQueryParsingToken> target = new ArrayList<>(selectionListTokens.size());
-		for (int i = 0; i < selectionListTokens.size(); i++) {
-			JpaQueryParsingToken token = selectionListTokens.get(i);
-			if (token.isA(TOKEN_AS)) {
-				i++;
-				continue;
-			}
-			target.add(token);
-		}
-		selectionListTokens = target;
-		return selectionListTokens;
-	}
 }

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

@@ -131,6 +131,12 @@ boolean getSpace() {
 		return this.space;
 	}
 
+	/**
+	 * Compare whether the given {@link JpaQueryParsingToken token} is equal to the one held by this instance.
+	 *
+	 * @param token must not be {@literal null}.
+	 * @return {@literal true} if both tokens are equals (using case-insensitive comparison).
+	 */
 	boolean isA(JpaQueryParsingToken token) {
 		return token.getToken().equalsIgnoreCase(this.getToken());
 	}

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,60 @@
+/*
+ * 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
+ * @since 3.2.5
+ */
+class QueryTransformers {
+
+	/**
+	 * Filter a token list from a {@code SELECT} clause to be used within a count query. That is, filter any {@code AS …}
+	 * aliases.
+	 *
+	 * @param selection the input selection.
+	 * @return filtered selection to be used with count queries.
+	 */
+	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/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() {