Adds support for more SelectBody types in JSqlParserQueryEhancer.

We now support `ValuesStatement` and `SetOperationList`. This allows native queries to use `union`, `except`, and `with` statements in native SQL queries.

Closes #2578.

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

@@ -29,9 +29,13 @@
 import net.sf.jsqlparser.statement.select.OrderByElement;
 import net.sf.jsqlparser.statement.select.PlainSelect;
 import net.sf.jsqlparser.statement.select.Select;
+import net.sf.jsqlparser.statement.select.SelectBody;
 import net.sf.jsqlparser.statement.select.SelectExpressionItem;
 import net.sf.jsqlparser.statement.select.SelectItem;
+import net.sf.jsqlparser.statement.select.SetOperationList;
+import net.sf.jsqlparser.statement.select.WithItem;
 import net.sf.jsqlparser.statement.update.Update;
+import net.sf.jsqlparser.statement.values.ValuesStatement;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -107,6 +111,13 @@ public String applySorting(Sort sort, @Nullable String alias) {
  }
 
  Select selectStatement = parseSelectStatement(queryString);
+
+ if (selectStatement.getSelectBody()instanceof SetOperationList setOperationList) {
+ return applySortingToSetOperationList(setOperationList, sort);
+ } else if (!(selectStatement.getSelectBody() instanceof PlainSelect)) {
+ return queryString;
+ }
+
  PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
 
  final Set<String> joinAliases = getJoinAliases(selectBody);
@@ -115,7 +126,7 @@ public String applySorting(Sort sort, @Nullable String alias) {
 
  List<OrderByElement> orderByElements = sort.stream() //
  .map(order -> getOrderClause(joinAliases, selectionAliases, alias, order)) //
- .collect(Collectors.toList());
+ .toList();
 
  if (CollectionUtils.isEmpty(selectBody.getOrderByElements())) {
  selectBody.setOrderByElements(new ArrayList<>());
@@ -127,6 +138,33 @@ public String applySorting(Sort sort, @Nullable String alias) {
 
  }
 
+ /**
+ * Returns the {@link SetOperationList} as a string query with {@link Sort}s applied in the right order.
+ * 
+ * @param setOperationListStatement
+ * @param sort
+ * @return
+ */
+ private String applySortingToSetOperationList(SetOperationList setOperationListStatement, Sort sort) {
+
+ // special case: ValuesStatements are detected as nested OperationListStatements
+ if (setOperationListStatement.getSelects().stream().anyMatch(ValuesStatement.class::isInstance)) {
+ return setOperationListStatement.toString();
+ }
+
+ // if (CollectionUtils.isEmpty(setOperationListStatement.getOrderByElements())) {
+ if (setOperationListStatement.getOrderByElements() == null) {
+ setOperationListStatement.setOrderByElements(new ArrayList<>());
+ }
+
+ List<OrderByElement> orderByElements = sort.stream() //
+ .map(order -> getOrderClause(Collections.emptySet(), Collections.emptySet(), null, order)) //
+ .toList();
+ setOperationListStatement.getOrderByElements().addAll(orderByElements);
+
+ return setOperationListStatement.toString();
+ }
+
  /**
  * Returns the aliases used inside the selection part in the query.
  *
@@ -175,7 +213,12 @@ private Set<String> getJoinAliases(String query) {
  return new HashSet<>();
  }
 
- return getJoinAliases((PlainSelect) parseSelectStatement(query).getSelectBody());
+ Select selectStatement = parseSelectStatement(query);
+ if (selectStatement.getSelectBody()instanceof PlainSelect selectBody) {
+ return getJoinAliases(selectBody);
+ }
+
+ return new HashSet<>();
  }
 
  /**
@@ -259,6 +302,17 @@ private String detectAlias(String query) {
  }
 
  Select selectStatement = parseSelectStatement(query);
+
+ /*
+ For all the other types ({@link ValuesStatement} and {@link SetOperationList}) it does not make sense to provide 
+ alias since:
+ * ValuesStatement has no alias
+ * SetOperation can have multiple alias for each operation item
+ */
+ if (!(selectStatement.getSelectBody() instanceof PlainSelect)) {
+ return null;
+ }
+
  PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
  return detectAlias(selectBody);
  }
@@ -273,6 +327,10 @@ private String detectAlias(String query) {
  @Nullable
  private static String detectAlias(PlainSelect selectBody) {
 
+ if (selectBody.getFromItem() == null) {
+ return null;
+ }
+
  Alias alias = selectBody.getFromItem().getAlias();
  return alias == null ? null : alias.getName();
  }
@@ -287,6 +345,14 @@ public String createCountQueryFor(@Nullable String countProjection) {
  Assert.hasText(this.query.getQueryString(), "OriginalQuery must not be null or empty");
 
  Select selectStatement = parseSelectStatement(this.query.getQueryString());
+
+ /*
+ We only support count queries for {@link PlainSelect}.
+ */
+ if (!(selectStatement.getSelectBody() instanceof PlainSelect)) {
+ return this.query.getQueryString();
+ }
+
  PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
 
  // remove order by
@@ -322,8 +388,15 @@ public String createCountQueryFor(@Nullable String countProjection) {
  Function jSqlCount = getJSqlCount(Collections.singletonList(countProp), distinct);
  selectBody.setSelectItems(Collections.singletonList(new SelectExpressionItem(jSqlCount)));
 
- return selectBody.toString();
+ if (CollectionUtils.isEmpty(selectStatement.getWithItemsList())) {
+ return selectBody.toString();
+ }
 
+ String withStatements = selectStatement.getWithItemsList().stream() //
+ .map(WithItem::toString) //
+ .collect(Collectors.joining(","));
+
+ return "with " + withStatements + "\n" + selectBody;
  }
 
  @Override
@@ -336,9 +409,23 @@ public String getProjection() {
  Assert.hasText(query.getQueryString(), "Query must not be null or empty");
 
  Select selectStatement = parseSelectStatement(query.getQueryString());
- PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
 
- return selectBody.getSelectItems() //
+ if (selectStatement.getSelectBody() instanceof ValuesStatement) {
+ return "";
+ }
+
+ SelectBody selectBody = selectStatement.getSelectBody();
+
+ if (selectStatement.getSelectBody()instanceof SetOperationList setOperationList) {
+ // using the first one since for setoperations the projection has to be the same
+ selectBody = setOperationList.getSelects().get(0);
+
+ if (!(selectBody instanceof PlainSelect)) {
+ return "";
+ }
+ }
+
+ return ((PlainSelect) selectBody).getSelectItems() //
  .stream() //
  .map(Object::toString) //
  .collect(Collectors.joining(", ")).trim();

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

@@ -2902,6 +2902,72 @@ public void correctlyBuildSortClauseWhenSortingByFunctionAliasAndFunctionContain
  repository.findAllAndSortByFunctionResultNamedParameter("prefix", "suffix", Sort.by("idWithPrefixAndSuffix"));
  }
 
+ @Test // GH-2578
+ void simpleNativeExceptTest() {
+
+ flushTestUsers();
+
+ List<String> foundIds = repository.findWithSimpleExceptNative();
+
+ assertThat(foundIds) //
+ .isNotEmpty() //
+ .contains("Oliver", "kevin");
+ }
+
+ @Test // GH-2578
+ void simpleNativeUnionTest() {
+
+ flushTestUsers();
+
+ List<String> foundIds = repository.findWithSimpleUnionNative();
+
+ assertThat(foundIds) //
+ .isNotEmpty() //
+ .containsExactlyInAnyOrder("Dave", "Joachim", "Oliver", "kevin");
+ }
+
+ @Test // GH-2578
+ void complexNativeExceptTest() {
+
+ flushTestUsers();
+
+ List<String> foundIds = repository.findWithComplexExceptNative();
+
+ assertThat(foundIds).containsExactly("Oliver", "kevin");
+ }
+
+ @Test // GH-2578
+ void simpleValuesStatementNative() {
+
+ flushTestUsers();
+
+ List<Integer> foundIds = repository.valuesStatementNative();
+
+ assertThat(foundIds).containsExactly(1);
+ }
+
+ @Test // GH-2578
+ void withStatementNative() {
+
+ flushTestUsers();
+
+ List<User> foundData = repository.withNativeStatement();
+
+ assertThat(foundData) //
+ .map(User::getFirstname) //
+ .containsExactly("Joachim", "Dave", "kevin");
+ }
+
+ @Test // GH-2578
+ void complexWithNativeStatement() {
+
+ flushTestUsers();
+
+ List<String> foundData = repository.complexWithNativeStatement();
+
+ assertThat(foundData).containsExactly("joachim", "dave", "kevin");
+ }
+
  private Page<User> executeSpecWithSort(Sort sort) {
 
  flushTestUsers();