Detect alias for merge statement for JSqlParserQueryEnhancer.

Related: #2641.

Author: gderemetz

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

@@ -18,6 +18,20 @@
 import static org.springframework.data.jpa.repository.query.JSqlParserUtils.*;
 import static org.springframework.data.jpa.repository.query.QueryUtils.*;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import org.springframework.data.domain.Sort;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+import org.springframework.util.CollectionUtils;
+import org.springframework.util.StringUtils;
+
 import net.sf.jsqlparser.JSQLParserException;
 import net.sf.jsqlparser.expression.Alias;
 import net.sf.jsqlparser.expression.Expression;
@@ -26,8 +40,8 @@
 import net.sf.jsqlparser.schema.Column;
 import net.sf.jsqlparser.statement.Statement;
 import net.sf.jsqlparser.statement.delete.Delete;
-import net.sf.jsqlparser.statement.merge.Merge;
 import net.sf.jsqlparser.statement.insert.Insert;
+import net.sf.jsqlparser.statement.merge.Merge;
 import net.sf.jsqlparser.statement.select.OrderByElement;
 import net.sf.jsqlparser.statement.select.PlainSelect;
 import net.sf.jsqlparser.statement.select.Select;
@@ -39,20 +53,6 @@
 import net.sf.jsqlparser.statement.update.Update;
 import net.sf.jsqlparser.statement.values.ValuesStatement;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Objects;
-import java.util.Set;
-import java.util.stream.Collectors;
-
-import org.springframework.data.domain.Sort;
-import org.springframework.lang.Nullable;
-import org.springframework.util.Assert;
-import org.springframework.util.CollectionUtils;
-import org.springframework.util.StringUtils;
-
 /**
  * The implementation of {@link QueryEnhancer} using JSqlParser.
  *
@@ -304,24 +304,28 @@ public String detectAlias() {
 	@Nullable
 	private String detectAlias(String query) {
 
-		if (this.parsedType != ParsedType.SELECT) {
-			return null;
-		}
-
-		Select selectStatement = parseSelectStatement(query);
+		if (ParsedType.MERGE.equals(this.parsedType)) {
+			Merge mergeStatement = parseSelectStatement(query, Merge.class);
+			return detectAlias(mergeStatement);
+
+		} else if (ParsedType.SELECT.equals(this.parsedType)) {
+			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;
+			}
 
-		/*
-		  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);
 		}
 
-		PlainSelect selectBody = (PlainSelect) selectStatement.getSelectBody();
-		return detectAlias(selectBody);
+		return null;
 	}
 
 	/**
@@ -332,7 +336,7 @@ For all the other types ({@link ValuesStatement} and {@link SetOperationList}) i
 	 * @return Might return {@literal null}.
 	 */
 	@Nullable
-	private static String detectAlias(PlainSelect selectBody) {
+	private String detectAlias(PlainSelect selectBody) {
 
 		if (selectBody.getFromItem() == null) {
 			return null;
@@ -342,6 +346,18 @@ private static String detectAlias(PlainSelect selectBody) {
 		return alias == null ? null : alias.getName();
 	}
 
+	/**
+	 * Resolves the alias for the given {@link Merge} statement.
+	 *
+	 * @param mergeStatement must not be {@literal null}.
+	 * @return Might return {@literal null}.
+	 */
+	@Nullable
+	private String detectAlias(Merge mergeStatement) {
+		Alias alias = mergeStatement.getUsingAlias();
+		return alias == null ? null : alias.getName();
+	}
+
 	@Override
 	public String createCountQueryFor(@Nullable String countProjection) {
 
@@ -449,15 +465,25 @@ public Set<String> getJoinAliases() {
 	 * @param query the query to parse
 	 * @return the parsed query
 	 */
-	private static Select parseSelectStatement(String query) {
+	private <T extends Statement> T parseSelectStatement(String query, Class<T> classOfT) {
 
 		try {
-			return (Select) CCJSqlParserUtil.parse(query);
+			return classOfT.cast(CCJSqlParserUtil.parse(query));
 		} catch (JSQLParserException e) {
 			throw new IllegalArgumentException("The query you provided is not a valid SQL Query", e);
 		}
 	}
 
+	/**
+	 * Parses a query string with JSqlParser.
+	 *
+	 * @param query the query to parse
+	 * @return the parsed query
+	 */
+	private Select parseSelectStatement(String query) {
+		return parseSelectStatement(query, Select.class);
+	}
+
 	/**
 	 * Checks whether a given projection only contains a single column definition (aka without functions, etc.)
 	 *

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

@@ -921,14 +921,17 @@ void insertStatementIsProcessedSameAsDefault(String insertQuery) {
 		assertThat(queryEnhancer.hasConstructorExpression()).isFalse();
 	}
 
-	@Test // GH-2641
-	void mergeStatementWorksWithJSqlParser() {
-		String query = "merge into a using (select id, value from b) query on (a.id = query.id) when matched then update set a.value = value";
+	@ParameterizedTest // GH-2641
+	@MethodSource("mergeStatementWorksWithJSqlParserSource")
+	void mergeStatementWorksWithJSqlParser(String query, String alias) {
 		StringQuery stringQuery = new StringQuery(query, true);
 		QueryEnhancer queryEnhancer = QueryEnhancerFactory.forQuery(stringQuery);
 
+		assertThat(queryEnhancer.detectAlias()).isEqualTo(alias);
+		assertThat(QueryUtils.detectAlias(query)).isNull();
+
 		assertThat(queryEnhancer.getJoinAliases()).isEmpty();
-		assertThat(queryEnhancer.detectAlias()).isNull();
+		assertThat(queryEnhancer.detectAlias()).isEqualTo(alias);
 		assertThat(queryEnhancer.getProjection()).isEmpty();
 		assertThat(queryEnhancer.hasConstructorExpression()).isFalse();
 	}
@@ -940,6 +943,15 @@ public static Stream<Arguments> insertStatementIsProcessedSameAsDefaultSource()
 		);
 	}
 
+	public static Stream<Arguments> mergeStatementWorksWithJSqlParserSource() {
+		return Stream.of(Arguments.of(
+				"merge into a using (select id, value from b) query on (a.id = query.id) when matched then update set a.value = value",
+				"query"),
+				Arguments.of(
+						"merge into a using (select id2, value from b) on (id = id2) when matched then update set a.value = value",
+						null));
+	}
+
 	public static Stream<Arguments> detectsJoinAliasesCorrectlySource() {
 
 		return Stream.of( //