Make JSqlParserQueryEnhancer work with updating statements.

This implementation of QueryEnhancer was originally designed for SELECT statements. This commit now handles DELETE and UPDATE operations by side-stepping any sorting or other changes.

Keep in mind that "enhancing" non selects does not have any effect on them (and the current default implementation `QueryUtils` does not care either aka it often just returns the same query, null or empty string).

Closes #2555

Author: DiegoKrupitza

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

@@ -24,11 +24,14 @@
 import net.sf.jsqlparser.expression.Function;
 import net.sf.jsqlparser.parser.CCJSqlParserUtil;
 import net.sf.jsqlparser.schema.Column;
+import net.sf.jsqlparser.statement.Statement;
+import net.sf.jsqlparser.statement.delete.Delete;
 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.SelectExpressionItem;
 import net.sf.jsqlparser.statement.select.SelectItem;
+import net.sf.jsqlparser.statement.update.Update;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -54,20 +57,49 @@
 public class JSqlParserQueryEnhancer implements QueryEnhancer {
 
 	private final DeclaredQuery query;
+	private final ParsedType parsedType;
 
 	/**
 	 * @param query the query we want to enhance. Must not be {@literal null}.
 	 */
 	public JSqlParserQueryEnhancer(DeclaredQuery query) {
 		this.query = query;
+		this.parsedType = detectParsedType();
+	}
+
+	/**
+	 * Detects what type of query is provided.
+	 * 
+	 * @return the parsed type
+	 */
+	private ParsedType detectParsedType() {
+		try {
+			Statement statement = CCJSqlParserUtil.parse(this.query.getQueryString());
+
+			if (statement instanceof Update) {
+				return ParsedType.UPDATE;
+			} else if (statement instanceof Delete) {
+				return ParsedType.DELETE;
+			} else if (statement instanceof Select) {
+				return ParsedType.SELECT;
+			} else {
+				return ParsedType.SELECT;
+			}
+
+		} catch (JSQLParserException e) {
+			throw new IllegalArgumentException("The query you provided is not a valid SQL Query!", e);
+		}
 	}
 
 	@Override
 	public String applySorting(Sort sort, @Nullable String alias) {
-
 		String queryString = query.getQueryString();
 		Assert.hasText(queryString, "Query must not be null or empty!");
 
+		if (this.parsedType != ParsedType.SELECT) {
+			return queryString;
+		}
+
 		if (sort.isUnsorted()) {
 			return queryString;
 		}
@@ -120,6 +152,10 @@ private Set<String> getSelectionAliases(PlainSelect selectBody) {
 	 */
 	Set<String> getSelectionAliases() {
 
+		if (this.parsedType != ParsedType.SELECT) {
+			return new HashSet<>();
+		}
+
 		Select selectStatement = parseSelectStatement(this.query.getQueryString());
 		PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
 		return this.getSelectionAliases(selectBody);
@@ -132,6 +168,9 @@ Set<String> getSelectionAliases() {
 	 * @return a {@literal Set} of aliases used in the query. Guaranteed to be not {@literal null}.
 	 */
 	private Set<String> getJoinAliases(String query) {
+		if (this.parsedType != ParsedType.SELECT) {
+			return new HashSet<>();
+		}
 		return getJoinAliases((PlainSelect) parseSelectStatement(query).getSelectBody());
 	}
 
@@ -211,6 +250,10 @@ public String detectAlias() {
 	@Nullable
 	private String detectAlias(String query) {
 
+		if (this.parsedType != ParsedType.SELECT) {
+			return null;
+		}
+
 		Select selectStatement = parseSelectStatement(query);
 		PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
 		return detectAlias(selectBody);
@@ -233,6 +276,10 @@ private static String detectAlias(PlainSelect selectBody) {
 	@Override
 	public String createCountQueryFor(@Nullable String countProjection) {
 
+		if (this.parsedType != ParsedType.SELECT) {
+			return this.query.getQueryString();
+		}
+
 		Assert.hasText(this.query.getQueryString(), "OriginalQuery must not be null or empty!");
 
 		Select selectStatement = parseSelectStatement(this.query.getQueryString());
@@ -278,6 +325,10 @@ public String createCountQueryFor(@Nullable String countProjection) {
 	@Override
 	public String getProjection() {
 
+		if (this.parsedType != ParsedType.SELECT) {
+			return "";
+		}
+
 		Assert.hasText(query.getQueryString(), "Query must not be null or empty!");
 
 		Select selectStatement = parseSelectStatement(query.getQueryString());
@@ -327,3 +378,15 @@ public DeclaredQuery getQuery() {
 		return this.query;
 	}
 }
+
+/**
+ * An enum to represent the top level parsed statement of the provided query.
+ * <ul>
+ * <li>{@code ParsedType.DELETE}: means the top level statement is {@link Delete}</li>
+ * <li>{@code ParsedType.UPDATE}: means the top level statement is {@link Update}</li>
+ * <li>{@code ParsedType.SELECT}: means the top level statement is {@link Select}</li>
+ * </ul>
+ */
+enum ParsedType {
+	DELETE, UPDATE, SELECT;
+}

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

@@ -2865,6 +2865,19 @@ void existsWithSpec() {
 		assertThat(repository.exists(hundredYearsOld)).isTrue();
 	}
 
+	@Test // GH-2555
+	void modifyingUpdateNativeQueryWorksWithJSQLParser() {
+		flushTestUsers();
+
+		Optional<User> byIdUser = repository.findById(firstUser.getId());
+		assertThat(byIdUser).isPresent().map(User::isActive).get().isEqualTo(true);
+
+		repository.setActiveToFalseWithModifyingNative(byIdUser.get().getId());
+
+		Optional<User> afterUpdate = repository.findById(firstUser.getId());
+		assertThat(afterUpdate).isPresent().map(User::isActive).get().isEqualTo(false);
+	}
+
 	@Test // GH-1262
 	void deleteWithSpec() {
 

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

@@ -731,6 +731,25 @@ void countQueryUsesCorrectVariable() {
 		assertThat(countQueryFor).isEqualTo("SELECT count(test) FROM (SELECT * FROM test) AS test");
 	}
 
+	@Test // GH-2555
+	void modifyingQueriesAreDetectedCorrectly() {
+		String modifyingQuery = "update userinfo user set user.is_in_treatment = false where user.id = :userId";
+
+		String aliasNotConsideringQueryType = QueryUtils.detectAlias(modifyingQuery);
+		String projectionNotConsideringQueryType = QueryUtils.getProjection(modifyingQuery);
+		boolean constructorExpressionNotConsideringQueryType = QueryUtils.hasConstructorExpression(modifyingQuery);
+		String countQueryForNotConsiderQueryType = QueryUtils.createCountQueryFor(modifyingQuery);
+
+		StringQuery modiQuery = new StringQuery(modifyingQuery, true);
+
+		assertThat(modiQuery.getAlias()).isEqualToIgnoringCase(aliasNotConsideringQueryType);
+		assertThat(modiQuery.getProjection()).isEqualToIgnoringCase(projectionNotConsideringQueryType);
+		assertThat(modiQuery.hasConstructorExpression()).isEqualTo(constructorExpressionNotConsideringQueryType);
+
+		assertThat(countQueryForNotConsiderQueryType).isEqualToIgnoringCase(modifyingQuery);
+		assertThat(QueryEnhancerFactory.forQuery(modiQuery).createCountQueryFor()).isEqualToIgnoringCase(modifyingQuery);
+	}
+
 	public static Stream<Arguments> detectsJoinAliasesCorrectlySource() {
 
 		return Stream.of( //

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

@@ -637,6 +637,11 @@ List<String> findAllAndSortByFunctionResultPositionalParameter(
 	List<String> findAllAndSortByFunctionResultNamedParameter(@Param("namedParameter1") String namedParameter1,
 			@Param("namedParameter2") String namedParameter2, Sort sort);
 
+	// GH-2555
+	@Modifying(clearAutomatically = true)
+	@Query(value = "update SD_User u set u.active = false where u.id = :userId", nativeQuery = true)
+	void setActiveToFalseWithModifyingNative(@Param("userId") int userId);
+
 	interface RolesAndFirstname {
 
 		String getFirstname();