Drop AS token when creating JPQL count query using DISTINCT.

Fixes: #3269
Original pull request: #3276

Author: christophstrobl

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 = getCountSelection(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,19 @@ public List<JpaQueryParsingToken> visitInstantiation(HqlParser.InstantiationCont
 	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,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/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);
 	}