Fix order by rendering for queries containing UNION.

Make sure to append space after order by clause and fix alias detection for wrapped sub select.
Also make sure to ignore alias used in subselect so they do not conflict with root ones.

Render order by only on full select if set operator is present in EQL.

Closes #3630
Original pull request: #3429
See #3427

Author: christophstrobl

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

@@ -24,6 +24,7 @@
 import org.springframework.data.domain.Sort;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.ObjectUtils;
 
 /**
  * An ANTLR {@link org.antlr.v4.runtime.tree.ParseTreeVisitor} that transforms a parsed EQL query.
@@ -105,27 +106,39 @@ public List<JpaQueryParsingToken> visitSelect_statement(EqlParser.Select_stateme
 
 		if (!countQuery) {
 
-			if (ctx.orderby_clause() != null) {
-				tokens.addAll(visit(ctx.orderby_clause()));
+			doVisitOrderBy(tokens, ctx, ObjectUtils.isEmpty(ctx.setOperator()) ? this.sort : Sort.unsorted());
+
+			for (int i = 0; i < ctx.setOperator().size(); i++) {
+
+				tokens.addAll(visit(ctx.setOperator(i)));
+				tokens.addAll(visit(ctx.select_statement(i)));
 			}
 
-			if (sort.isSorted()) {
+		}
 
-				if (ctx.orderby_clause() != null) {
+		return tokens;
+	}
 
-					NOSPACE(tokens);
-					tokens.add(TOKEN_COMMA);
-				} else {
+	private void doVisitOrderBy(List<JpaQueryParsingToken> tokens, EqlParser.Select_statementContext ctx, Sort sort) {
 
-					SPACE(tokens);
-					tokens.add(TOKEN_ORDER_BY);
-				}
+		if (ctx.orderby_clause() != null) {
+			tokens.addAll(visit(ctx.orderby_clause()));
+		}
 
-				tokens.addAll(transformerSupport.generateOrderByArguments(primaryFromAlias, sort));
+		if (sort.isSorted()) {
+
+			if (ctx.orderby_clause() != null) {
+
+				NOSPACE(tokens);
+				tokens.add(TOKEN_COMMA);
+			} else {
+
+				SPACE(tokens);
+				tokens.add(TOKEN_ORDER_BY);
 			}
-		}
 
-		return tokens;
+			tokens.addAll(transformerSupport.generateOrderByArguments(primaryFromAlias, sort));
+		}
 	}
 
 	@Override

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

@@ -26,6 +26,7 @@
 import org.springframework.data.jpa.repository.query.HqlParser.SelectionContext;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.ObjectUtils;
 
 /**
  * An ANTLR {@link org.antlr.v4.runtime.tree.ParseTreeVisitor} that transforms a parsed HQL query.
@@ -105,8 +106,7 @@ private static boolean isSubquery(ParserRuleContext ctx) {
 		}
 	}
 
-	@Override
-	public List<JpaQueryParsingToken> visitOrderedQuery(HqlParser.OrderedQueryContext ctx) {
+	private List<JpaQueryParsingToken> visitOrderedQuery(HqlParser.OrderedQueryContext ctx, Sort sort) {
 
 		List<JpaQueryParsingToken> tokens = newArrayList();
 
@@ -190,6 +190,40 @@ public List<JpaQueryParsingToken> visitFromQuery(HqlParser.FromQueryContext ctx)
 		return tokens;
 	}
 
+	@Override
+	public List<JpaQueryParsingToken> visitQueryExpression(HqlParser.QueryExpressionContext ctx) {
+
+		if (ObjectUtils.isEmpty(ctx.setOperator())) {
+			return super.visitQueryExpression(ctx);
+		}
+
+		List<JpaQueryParsingToken> builder = new ArrayList<>();
+		if (ctx.withClause() != null) {
+			builder.addAll(visit(ctx.withClause()));
+		}
+
+		List<HqlParser.OrderedQueryContext> orderedQueries = ctx.orderedQuery();
+		for (int i = 0; i < orderedQueries.size(); i++) {
+
+			if (i != 0) {
+				builder.addAll(visit(ctx.setOperator(i - 1)));
+			}
+
+			if (i == orderedQueries.size() - 1) {
+				builder.addAll(visitOrderedQuery(ctx.orderedQuery(i), this.sort));
+			} else {
+				builder.addAll(visitOrderedQuery(ctx.orderedQuery(i), Sort.unsorted()));
+			}
+		}
+
+		return builder;
+	}
+
+	@Override
+	public List<JpaQueryParsingToken> visitOrderedQuery(HqlParser.OrderedQueryContext ctx) {
+		return visitOrderedQuery(ctx, this.sort);
+	}
+
 	@Override
 	public List<JpaQueryParsingToken> visitQueryOrder(HqlParser.QueryOrderContext ctx) {
 
@@ -325,7 +359,7 @@ public List<JpaQueryParsingToken> visitVariable(HqlParser.VariableContext ctx) {
 
 		List<JpaQueryParsingToken> tokens = super.visitVariable(ctx);
 
-		if (ctx.identifier() != null) {
+		if (ctx.identifier() != null && !tokens.isEmpty() && !isSubquery(ctx)) {
 			transformerSupport.registerAlias(tokens.get(tokens.size() - 1).getToken());
 		}
 
@@ -335,7 +369,7 @@ public List<JpaQueryParsingToken> visitVariable(HqlParser.VariableContext ctx) {
 	@Override
 	public List<JpaQueryParsingToken> visitSelection(SelectionContext ctx) {
 
-		if(!countQuery || isSubquery(ctx)) {
+		if (!countQuery || isSubquery(ctx)) {
 			return super.visitSelection(ctx);
 		}
 

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

@@ -16,7 +16,7 @@
 
 /**
  * Transformational operations needed to support either {@link HqlQueryTransformer} or {@link JpqlQueryTransformer}.
- * 
+ *
  * @author Greg Turnquist
  * @author Donghun Shin
  * @since 3.1
@@ -47,12 +47,12 @@ void registerAlias(String token) {
 	/**
 	 * Using the primary {@literal FROM} clause's alias and a {@link Sort}, construct all the {@literal ORDER BY}
 	 * arguments.
-	 * 
+	 *
 	 * @param primaryFromAlias
 	 * @param sort
 	 * @return
 	 */
-	List<JpaQueryParsingToken> generateOrderByArguments(String primaryFromAlias, Sort sort) {
+	List<JpaQueryParsingToken> generateOrderByArguments(@Nullable String primaryFromAlias, Sort sort) {
 
 		List<JpaQueryParsingToken> tokens = new ArrayList<>();
 
@@ -98,7 +98,7 @@ private void checkSortExpression(Sort.Order order) {
 	/**
 	 * Using the {@code primaryFromAlias} and the {@link org.springframework.data.domain.Sort.Order}, construct a suitable
 	 * argument to be added to an {@literal ORDER BY} expression.
-	 * 
+	 *
 	 * @param primaryFromAlias
 	 * @param order
 	 * @return

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

@@ -753,6 +753,15 @@ void sortingRecognizesJoinAliases() {
 				""");
 	}
 
+	@Test // GH-3427
+	void sortShouldBeAppendedToFullSelectOnlyInCaseOfSetOperator() {
+
+		String source = "SELECT tb FROM Test tb WHERE (tb.type='A') UNION SELECT tb FROM Test tb WHERE (tb.type='B')";
+		String target = createQueryFor(source, Sort.by("Type").ascending());
+
+		assertThat(target).isEqualTo("SELECT tb FROM Test tb WHERE (tb.type = 'A') UNION SELECT tb FROM Test tb WHERE (tb.type = 'B') order by tb.Type asc");
+	}
+
 	static Stream<Arguments> queriesWithReservedWordsAsIdentifiers() {
 
 		return Stream.of( //

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

@@ -17,18 +17,22 @@
 
 import static org.assertj.core.api.Assertions.*;
 
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import java.util.stream.Stream;
 
 import org.assertj.core.api.SoftAssertions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.data.domain.PageRequest;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.jpa.domain.JpaSort;
 import org.springframework.lang.Nullable;
+import org.springframework.util.StringUtils;
 
 /**
  * Verify that HQL queries are properly transformed through the {@link JpaQueryEnhancer} and the {@link HqlQueryParser}.
@@ -1060,6 +1064,40 @@ void createsCountQueryUsingAliasCorrectly() {
 		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");
 	}
 
+	@Test // GH-3427
+	void sortShouldBeAppendedWithSpacingInCaseOfSetOperator() {
+
+        String source = "SELECT tb FROM Test tb WHERE (tb.type='A') UNION SELECT tb FROM Test tb WHERE (tb.type='B')";
+		String target = createQueryFor(source, Sort.by("Type").ascending());
+		
+		assertThat(target).isEqualTo("SELECT tb FROM Test tb WHERE (tb.type = 'A') UNION SELECT tb FROM Test tb WHERE (tb.type = 'B') order by tb.Type asc");
+	}
+
+	@ParameterizedTest // GH-3427
+	@ValueSource(strings = {"", "res"})
+	void sortShouldBeAppendedToSubSelectWithSetOperatorInSubselect(String alias) {
+
+		String prefix = StringUtils.hasText(alias) ? (alias + ".") : "";
+		String source = "SELECT %sname FROM (SELECT c.name as name FROM Category c UNION SELECT t.name as name FROM Tag t)".formatted(prefix);
+		if(StringUtils.hasText(alias)) {
+			source = source + " %s".formatted(alias);
+		}
+
+		String target = createQueryFor(source, Sort.by("name").ascending());
+
+		assertThat(target).contains(" UNION SELECT ").doesNotContainPattern(Pattern.compile(".*\\SUNION"));
+		assertThat(target).endsWith("order by %sname asc".formatted(prefix)).satisfies(it -> {
+			Pattern pattern = Pattern.compile("order by %sname".formatted(prefix));
+			Matcher matcher = pattern.matcher(target);
+			int count = 0;
+			while(matcher.find()) {
+				count++;
+			}
+			assertThat(count).describedAs("Found order by clause more than once in: \n%s", it).isOne();
+		});
+
+	}
+
 	private void assertCountQuery(String originalQuery, String countQuery) {
 		assertThat(createCountQueryFor(originalQuery)).isEqualTo(countQuery);
 	}