Fix query execution mode detection for aggregate types that implement Streamable.

odrotbohm · odrotbohm · commit 5ae27300a99a · 2023-07-01T01:56:45.000+02:00
We now short-circuit the QueryMethod.isCollectionQuery() algorithm in case we find the concrete domain type or any subclass of it. Fixes #2869.
diff --git a/src/main/java/org/springframework/data/repository/core/support/AbstractRepositoryMetadata.java b/src/main/java/org/springframework/data/repository/core/support/AbstractRepositoryMetadata.java
@@ -105,12 +105,13 @@ public TypeInformation<?> getReturnType(Method method) {
 	 * (non-Javadoc)
 	 * @see org.springframework.data.repository.core.RepositoryMetadata#getReturnedDomainClass(java.lang.reflect.Method)
 	 */
+	@Override
 	public Class<?> getReturnedDomainClass(Method method) {
 
 		TypeInformation<?> returnType = getReturnType(method);
+		returnType = ReactiveWrapperConverters.unwrapWrapperTypes(returnType);
 
-		return QueryExecutionConverters.unwrapWrapperTypes(ReactiveWrapperConverters.unwrapWrapperTypes(returnType))
-				.getType();
+		return QueryExecutionConverters.unwrapWrapperTypes(returnType, getDomainTypeInformation()).getType();
 	}
 
 	/*
diff --git a/src/main/java/org/springframework/data/repository/query/QueryMethod.java b/src/main/java/org/springframework/data/repository/query/QueryMethod.java
@@ -32,6 +32,7 @@
 import org.springframework.data.repository.util.ReactiveWrapperConverters;
 import org.springframework.data.util.ClassTypeInformation;
 import org.springframework.data.util.Lazy;
+import org.springframework.data.util.NullableWrapperConverters;
 import org.springframework.data.util.TypeInformation;
 import org.springframework.util.Assert;
 
@@ -265,7 +266,15 @@ private boolean calculateIsCollectionQuery() {
 			return false;
 		}
 
-		Class<?> returnType = metadata.getReturnType(method).getType();
+		TypeInformation<?> returnTypeInformation = metadata.getReturnType(method);
+
+		// Check against simple wrapper types first
+		if (metadata.getDomainTypeInformation()
+				.isAssignableFrom(NullableWrapperConverters.unwrapActualType(returnTypeInformation))) {
+			return false;
+		}
+
+		Class<?> returnType = returnTypeInformation.getType();
 
 		if (QueryExecutionConverters.supports(returnType) && !QueryExecutionConverters.isSingleValue(returnType)) {
 			return true;
diff --git a/src/main/java/org/springframework/data/repository/util/QueryExecutionConverters.java b/src/main/java/org/springframework/data/repository/util/QueryExecutionConverters.java
@@ -39,6 +39,7 @@
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.Slice;
 import org.springframework.data.geo.GeoResults;
+import org.springframework.data.util.ClassTypeInformation;
 import org.springframework.data.util.CustomCollections;
 import org.springframework.data.util.NullableWrapper;
 import org.springframework.data.util.NullableWrapperConverters;
@@ -86,6 +87,7 @@ public abstract class QueryExecutionConverters {
 	private static final Set<Class<?>> ALLOWED_PAGEABLE_TYPES = new HashSet<>();
 	private static final Map<Class<?>, ExecutionAdapter> EXECUTION_ADAPTER = new HashMap<>();
 	private static final Map<Class<?>, Boolean> supportsCache = new ConcurrentReferenceHashMap<>();
+	private static final TypeInformation<Void> VOID_INFORMATION = ClassTypeInformation.from(Void.class);
 
 	static {
 
@@ -235,15 +237,21 @@ public static Object unwrap(@Nullable Object source) {
 	}
 
 	/**
-	 * Recursively unwraps well known wrapper types from the given {@link TypeInformation}.
+	 * Recursively unwraps well known wrapper types from the given {@link TypeInformation} but aborts at the given
+	 * reference type.
 	 *
 	 * @param type must not be {@literal null}.
+	 * @param reference must not be {@literal null}.
 	 * @return will never be {@literal null}.
 	 */
-	public static TypeInformation<?> unwrapWrapperTypes(TypeInformation<?> type) {
+	public static TypeInformation<?> unwrapWrapperTypes(TypeInformation<?> type, TypeInformation<?> reference) {
 
 		Assert.notNull(type, "type must not be null");
 
+		if (reference.isAssignableFrom(type)) {
+			return type;
+		}
+
 		Class<?> rawType = type.getType();
 
 		boolean needToUnwrap = type.isCollectionLike() //
@@ -253,7 +261,17 @@ public static TypeInformation<?> unwrapWrapperTypes(TypeInformation<?> type) {
 				|| supports(rawType) //
 				|| Stream.class.isAssignableFrom(rawType);
 
-		return needToUnwrap ? unwrapWrapperTypes(type.getRequiredComponentType()) : type;
+		return needToUnwrap ? unwrapWrapperTypes(type.getRequiredComponentType(), reference) : type;
+	}
+
+	/**
+	 * Recursively unwraps well known wrapper types from the given {@link TypeInformation}.
+	 *
+	 * @param type must not be {@literal null}.
+	 * @return will never be {@literal null}.
+	 */
+	public static TypeInformation<?> unwrapWrapperTypes(TypeInformation<?> type) {
+		return unwrapWrapperTypes(type, VOID_INFORMATION);
 	}
 
 	/**
diff --git a/src/test/java/org/springframework/data/repository/core/support/AbstractRepositoryMetadataUnitTests.java b/src/test/java/org/springframework/data/repository/core/support/AbstractRepositoryMetadataUnitTests.java
@@ -21,15 +21,20 @@
 import java.lang.reflect.Method;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Optional;
+import java.util.stream.Stream;
 
+import org.junit.jupiter.api.DynamicTest;
 import org.junit.jupiter.api.Test;
-import org.springframework.core.ResolvableType;
+import org.junit.jupiter.api.TestFactory;
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.querydsl.User;
 import org.springframework.data.repository.PagingAndSortingRepository;
 import org.springframework.data.repository.Repository;
 import org.springframework.data.repository.core.RepositoryMetadata;
+import org.springframework.data.util.Streamable;
 
 /**
  * Unit tests for {@link AbstractRepositoryMetadata}.
@@ -112,6 +117,25 @@ void doesNotUnwrapCustomTypeImplementingIterable() throws Exception {
 		assertThat(metadata.getReturnedDomainClass(method)).isEqualTo(Container.class);
 	}
 
+	@TestFactory // GH-2869
+	Stream<DynamicTest> detectsReturnTypesForStreamableAggregates() throws Exception {
+
+		RepositoryMetadata metadata = AbstractRepositoryMetadata.getMetadata(StreamableAggregateRepository.class);
+		Stream<Entry<String, Class<?>>> methods = Stream.of(
+				Map.entry("findBy", StreamableAggregate.class),
+				Map.entry("findSubTypeBy", StreamableAggregateSubType.class),
+				Map.entry("findAllBy", StreamableAggregate.class),
+				Map.entry("findOptional", StreamableAggregate.class));
+
+		return DynamicTest.stream(methods, //
+				it -> it.getKey() + "'s returned domain class is " + it.getValue(), //
+				it -> {
+
+					Method method = StreamableAggregateRepository.class.getMethod(it.getKey());
+					assertThat(metadata.getReturnedDomainClass(method)).isEqualTo(it.getValue());
+				});
+	}
+
 	interface UserRepository extends Repository<User, Long> {
 
 		User findSingle();
@@ -153,4 +177,21 @@ abstract class Container implements Iterable<Element> {}
 	interface ContainerRepository extends Repository<Container, Long> {
 		Container someMethod();
 	}
+
+	// GH-2869
+
+	static abstract class StreamableAggregate implements Streamable<Object> {}
+
+	interface StreamableAggregateRepository extends Repository<StreamableAggregate, Object> {
+
+		StreamableAggregate findBy();
+
+		StreamableAggregateSubType findSubTypeBy();
+
+		Streamable<StreamableAggregate> findAllBy();
+
+		Optional<StreamableAggregate> findOptional();
+	}
+
+	static abstract class StreamableAggregateSubType extends StreamableAggregate {}
 }
diff --git a/src/test/java/org/springframework/data/repository/query/QueryMethodUnitTests.java b/src/test/java/org/springframework/data/repository/query/QueryMethodUnitTests.java
@@ -24,12 +24,17 @@
 import java.io.Serializable;
 import java.lang.reflect.Method;
 import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
 import java.util.stream.Stream;
 
 import org.eclipse.collections.api.list.ImmutableList;
+import org.junit.jupiter.api.DynamicTest;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestFactory;
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Slice;
@@ -39,6 +44,7 @@
 import org.springframework.data.repository.core.RepositoryMetadata;
 import org.springframework.data.repository.core.support.AbstractRepositoryMetadata;
 import org.springframework.data.repository.core.support.DefaultRepositoryMetadata;
+import org.springframework.data.util.Streamable;
 
 /**
  * Unit tests for {@link QueryMethod}.
@@ -258,6 +264,28 @@ void considersEclipseCollectionCollectionQuery() throws Exception {
 		assertThat(queryMethod.isCollectionQuery()).isTrue();
 	}
 
+	@TestFactory // GH-2869
+	Stream<DynamicTest> doesNotConsiderQueryMethodReturningAggregateImplementingStreamableACollectionQuery()
+			throws Exception {
+
+		RepositoryMetadata metadata = AbstractRepositoryMetadata.getMetadata(StreamableAggregateRepository.class);
+		Stream<Entry<String, Boolean>> stream = Stream.of(
+				Map.entry("findBy", false),
+				Map.entry("findSubTypeBy", false),
+				Map.entry("findAllBy", true),
+				Map.entry("findOptionalBy", false));
+
+		return DynamicTest.stream(stream, //
+				it -> it.getKey() + " considered collection query -> " + it.getValue(), //
+				it -> {
+
+					Method method = StreamableAggregateRepository.class.getMethod(it.getKey());
+					QueryMethod queryMethod = new QueryMethod(method, metadata, factory);
+
+					assertThat(queryMethod.isCollectionQuery()).isEqualTo(it.getValue());
+				});
+	}
+
 	interface SampleRepository extends Repository<User, Serializable> {
 
 		String pagingMethodWithInvalidReturnType(Pageable pageable);
@@ -325,4 +353,21 @@ abstract class Container implements Iterable<Element> {}
 	interface ContainerRepository extends Repository<Container, Long> {
 		Container someMethod();
 	}
+
+	// GH-2869
+
+	static abstract class StreamableAggregate implements Streamable<Object> {}
+
+	interface StreamableAggregateRepository extends Repository<StreamableAggregate, Object> {
+
+		StreamableAggregate findBy();
+
+		StreamableAggregateSubType findSubTypeBy();
+
+		Optional<StreamableAggregate> findOptionalBy();
+
+		Streamable<StreamableAggregate> findAllBy();
+	}
+
+	static abstract class StreamableAggregateSubType extends StreamableAggregate {}
 }
