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.

Author: DiegoKrupitza

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,14 @@ public String applySorting(Sort sort, @Nullable String alias) {
 		}
 
 		Select selectStatement = parseSelectStatement(queryString);
+
+		if (selectStatement.getSelectBody() instanceof SetOperationList) {
+			SetOperationList setOperationList = (SetOperationList) selectStatement.getSelectBody();
+			return applySortingToSetOperationList(setOperationList, sort);
+		} else if (!(selectStatement.getSelectBody() instanceof PlainSelect)) {
+			return queryString;
+		}
+
 		PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
 
 		final Set<String> joinAliases = getJoinAliases(selectBody);
@@ -127,6 +139,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)) //
+				.collect(Collectors.toList());
+		setOperationListStatement.getOrderByElements().addAll(orderByElements);
+
+		return setOperationListStatement.toString();
+	}
+
 	/**
 	 * Returns the aliases used inside the selection part in the query.
 	 *
@@ -175,7 +214,13 @@ private Set<String> getJoinAliases(String query) {
 			return new HashSet<>();
 		}
 
-		return getJoinAliases((PlainSelect) parseSelectStatement(query).getSelectBody());
+		Select selectStatement = parseSelectStatement(query);
+		if (selectStatement.getSelectBody() instanceof PlainSelect) {
+			PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
+			return getJoinAliases(selectBody);
+		}
+
+		return new HashSet<>();
 	}
 
 	/**
@@ -259,6 +304,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 +329,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 +347,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 +390,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 +411,24 @@ 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 setOperationList = (SetOperationList) selectStatement.getSelectBody();
+			// 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();

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

@@ -2702,6 +2702,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();