Detect alias for merge statement for JSqlParserQueryEnhancer.

Related: #2641.

Author: gderemetz

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

@@ -307,24 +307,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;
 	}
 
 	/**
@@ -335,7 +339,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;
@@ -345,6 +349,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) {
 
@@ -453,15 +469,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.)
 	 *

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

@@ -923,14 +923,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();
 	}
@@ -942,6 +945,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( //