DATAJPA-1827 - Consider wrapper types for Modifying JPA Query Execution

pom.xml

@@ -1,11 +1,13 @@
 <?xml version="1.0" encoding="UTF-8" standalone="no"?>
-<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
 
 	<modelVersion>4.0.0</modelVersion>
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-jpa</artifactId>
-	<version>2.5.0-SNAPSHOT</version>
+	<version>2.5.0-DATAJPA-1827-SNAPSHOT</version>
 
 	<name>Spring Data JPA</name>
 	<description>Spring Data module for JPA repositories.</description>
@@ -24,6 +26,7 @@
 		<eclipselink>2.7.5</eclipselink>
 		<hibernate>5.4.8.Final</hibernate>
 		<mockito>2.19.1</mockito>
+		<vavr>0.10.3</vavr>
 		<hibernate.groupId>org.hibernate</hibernate.groupId>
 		<springdata.commons>2.5.0-SNAPSHOT</springdata.commons>
 
@@ -197,6 +200,13 @@
 			<optional>true</optional>
 		</dependency>
 
+		<dependency>
+			<groupId>io.vavr</groupId>
+			<artifactId>vavr</artifactId>
+			<version>${vavr}</version>
+			<scope>test</scope>
+		</dependency>
+
 		<!-- Persistence providers -->
 
 		<dependency>

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

@@ -97,7 +97,7 @@ public Object execute(AbstractJpaQuery query, JpaParametersParameterAccessor acc
 		JpaQueryMethod queryMethod = query.getQueryMethod();
 		Class<?> requiredType = queryMethod.getReturnType();
 
-		if (void.class.equals(requiredType) || requiredType.isAssignableFrom(result.getClass())) {
+		if (ClassUtils.isAssignable(requiredType, void.class) || ClassUtils.isAssignableValue(requiredType, result)) {
 			return result;
 		}
 
@@ -218,10 +218,11 @@ public ModifyingExecution(JpaQueryMethod method, EntityManager em) {
 
 			Class<?> returnType = method.getReturnType();
 
-			boolean isVoid = void.class.equals(returnType) || Void.class.equals(returnType);
-			boolean isInt = int.class.equals(returnType) || Integer.class.equals(returnType);
+			boolean isVoid = ClassUtils.isAssignable(returnType, Void.class);
+			boolean isInt = ClassUtils.isAssignable(returnType, Integer.class);
 
-			Assert.isTrue(isInt || isVoid, "Modifying queries can only use void or int/Integer as return type!");
+			Assert.isTrue(isInt || isVoid,
+					"Modifying queries can only use void or int/Integer as return type! Offending method: " + method);
 
 			this.em = em;
 			this.flush = method.getFlushAutomatically();

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

@@ -40,7 +40,9 @@
 import org.springframework.data.repository.query.Parameter;
 import org.springframework.data.repository.query.Parameters;
 import org.springframework.data.repository.query.QueryMethod;
+import org.springframework.data.repository.util.QueryExecutionConverters;
 import org.springframework.data.util.Lazy;
+import org.springframework.data.util.TypeInformation;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
@@ -79,6 +81,7 @@ public class JpaQueryMethod extends QueryMethod {
 
 	private final QueryExtractor extractor;
 	private final Method method;
+	private final Class<?> returnType;
 
 	private @Nullable StoredProcedureAttributes storedProcedureAttributes;
 	private final Lazy<LockModeType> lockModeType;
@@ -107,6 +110,7 @@ protected JpaQueryMethod(Method method, RepositoryMetadata metadata, ProjectionF
 		Assert.notNull(extractor, "Query extractor must not be null!");
 
 		this.method = method;
+		this.returnType = potentiallyUnwrapReturnTypeFor(metadata, method);
 		this.extractor = extractor;
 		this.lockModeType = Lazy
 				.of(() -> (LockModeType) Optional.ofNullable(AnnotatedElementUtils.findMergedAnnotation(method, Lock.class)) //
@@ -126,8 +130,7 @@ protected JpaQueryMethod(Method method, RepositoryMetadata metadata, ProjectionF
 			return new JpaEntityGraph(entityGraph, getNamedQueryName());
 		});
 		this.isNativeQuery = Lazy.of(() -> getAnnotationValue("nativeQuery", Boolean.class));
-		this.isCollectionQuery = Lazy
-				.of(() -> super.isCollectionQuery() && !NATIVE_ARRAY_TYPES.contains(method.getReturnType()));
+		this.isCollectionQuery = Lazy.of(() -> super.isCollectionQuery() && !NATIVE_ARRAY_TYPES.contains(this.returnType));
 		this.isProcedureQuery = Lazy.of(() -> AnnotationUtils.findAnnotation(method, Procedure.class) != null);
 		this.entityMetadata = Lazy.of(() -> new DefaultJpaEntityMetadata<>(getDomainClass()));
 
@@ -136,6 +139,18 @@ protected JpaQueryMethod(Method method, RepositoryMetadata metadata, ProjectionF
 		assertParameterNamesInAnnotatedQuery();
 	}
 
+	private static Class<?> potentiallyUnwrapReturnTypeFor(RepositoryMetadata metadata, Method method) {
+
+		TypeInformation<?> returnType = metadata.getReturnType(method);
+
+		while (QueryExecutionConverters.supports(returnType.getType())
+				|| QueryExecutionConverters.supportsUnwrapping(returnType.getType())) {
+			returnType = returnType.getRequiredComponentType();
+		}
+
+		return returnType.getType();
+	}
+
 	private void assertParameterNamesInAnnotatedQuery() {
 
 		String annotatedQuery = getAnnotatedQuery();
@@ -243,7 +258,7 @@ QueryExtractor getQueryExtractor() {
 	 * @return
 	 */
 	Class<?> getReturnType() {
-		return method.getReturnType();
+		return returnType;
 	}
 
 	/**

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

@@ -19,6 +19,9 @@
 import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.*;
 
+import io.vavr.control.Try;
+
+import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Optional;
@@ -38,8 +41,13 @@
 
 import org.springframework.data.domain.PageRequest;
 import org.springframework.data.domain.Pageable;
+import org.springframework.data.jpa.provider.QueryExtractor;
+import org.springframework.data.jpa.repository.Modifying;
 import org.springframework.data.jpa.repository.query.JpaQueryExecution.ModifyingExecution;
 import org.springframework.data.jpa.repository.query.JpaQueryExecution.PagedExecution;
+import org.springframework.data.projection.SpelAwareProxyProjectionFactory;
+import org.springframework.data.repository.Repository;
+import org.springframework.data.repository.core.support.DefaultRepositoryMetadata;
 
 /**
  * Unit test for {@link JpaQueryExecution}.
@@ -85,6 +93,27 @@ void rejectsNullBinder() {
 		assertThatIllegalArgumentException().isThrownBy(() -> new StubQueryExecution().execute(jpaQuery, null));
 	}
 
+	@Test // DATAJPA-1827
+	void supportsModifyingResultsUsingWrappers() throws Exception {
+
+		Method method = VavrRepository.class.getMethod("updateUsingVavrMethod");
+		DefaultRepositoryMetadata repositoryMetadata = new DefaultRepositoryMetadata(VavrRepository.class);
+		JpaQueryMethod queryMethod = new JpaQueryMethod(method, repositoryMetadata, new SpelAwareProxyProjectionFactory(),
+				mock(QueryExtractor.class));
+
+		new JpaQueryExecution.ModifyingExecution(queryMethod, mock(EntityManager.class));
+
+		assertThat(queryMethod.isModifyingQuery()).isTrue();
+	}
+
+	interface VavrRepository extends Repository<String, String> {
+
+		// Wrapped outcome allowed
+		@org.springframework.data.jpa.repository.Query("update Credential d set d.enabled = false where d.enabled = true")
+		@Modifying
+		Try<Integer> updateUsingVavrMethod();
+	}
+
 	@Test
 	void transformsNoResultExceptionToNull() {