Retain selection columns on query transformation without entity alias.

Closes #3744
spring-projects · Jan 14, 2025 · 6d13559 · 6d13559
1 parent f2d1a35
commit 6d13559
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 9 deletions.
diff --git a/...src/main/java/org/springframework/data/jpa/repository/query/HqlCountQueryTransformer.java b/...src/main/java/org/springframework/data/jpa/repository/query/HqlCountQueryTransformer.java
@@ -127,11 +127,6 @@ public QueryRendererBuilder visitFromRoot(HqlParser.FromRootContext ctx) {
 
 			if (ctx.variable() != null) {
 				builder.appendExpression(visit(ctx.variable()));
-
-			} else {
-
-				builder.append(TOKEN_AS);
-				builder.append(TOKEN_DOUBLE_UNDERSCORE);
 			}
 		} else if (ctx.subquery() != null) {
 
@@ -186,14 +181,26 @@ public QueryTokenStream visitSelectClause(HqlParser.SelectClauseContext ctx) {
 		boolean usesDistinct = ctx.DISTINCT() != null;
 		QueryRendererBuilder nested = QueryRenderer.builder();
 		if (countProjection == null) {
+			QueryTokenStream selection = visit(ctx.selectionList());
 			if (usesDistinct) {
 
 				nested.append(QueryTokens.expression(ctx.DISTINCT()));
-				nested.append(getDistinctCountSelection(visit(ctx.selectionList())));
+				nested.append(getDistinctCountSelection(selection));
 			} else {
 
 				// with CTE primary alias fails with hibernate (WITH entities AS (…) SELECT count(c) FROM entities c)
-				nested.append(containsCTE ? QueryTokens.token("*") : QueryTokens.token(primaryFromAlias));
+				if (containsCTE) {
+					nested.append(QueryTokens.token("*"));
+				} else {
+
+					if (selection.size() == 1) {
+						nested.append(selection);
+					} else if (primaryFromAlias != null) {
+						nested.append(QueryTokens.token(primaryFromAlias));
+					} else {
+						nested.append(QueryTokens.token("*"));
+					}
+				}
 			}
 		} else {
 			builder.append(QueryTokens.token(countProjection));
@@ -244,6 +251,7 @@ public QueryRendererBuilder visitQueryOrder(HqlParser.QueryOrderContext ctx) {
 	}
 
 	private QueryRendererBuilder visitSubQuerySelectClause(SelectClauseContext ctx, QueryRendererBuilder builder) {
+
 		if (ctx.DISTINCT() != null) {
 			builder.append(QueryTokens.expression(ctx.DISTINCT()));
 		}
@@ -258,8 +266,13 @@ private QueryRendererBuilder getDistinctCountSelection(QueryTokenStream selectio
 		CountSelectionTokenStream countSelection = CountSelectionTokenStream.create(selectionListbuilder);
 
 		if (countSelection.requiresPrimaryAlias()) {
-			// constructor
-			nested.append(QueryTokens.token(primaryFromAlias));
+
+			if (primaryFromAlias != null) {
+				// constructor
+				nested.append(QueryTokens.token(primaryFromAlias));
+			} else {
+				nested.append(countSelection.withoutConstructorExpression());
+			}
 		} else {
 			// keep all the select items to distinct against
 			nested.append(selectionListbuilder);

diff --git a/...pa/src/test/java/org/springframework/data/jpa/repository/query/QueryEnhancerTckTests.java b/...pa/src/test/java/org/springframework/data/jpa/repository/query/QueryEnhancerTckTests.java
@@ -68,6 +68,10 @@ static Stream<Arguments> nativeCountQueries() {
 						"select u from User as u", //
 						"select count(u) from User as u"),
 
+				Arguments.of( //
+						"SELECT id FROM Person", //
+						"select count(id) from Person"),
+
 				Arguments.of( //
 						"SELECT u FROM User u where u.foo.bar = ?", //
 						"select count(u) FROM User u where u.foo.bar = ?"),