Allow Custom Streamable return type.

Author: christophstrobl

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/aot/MongoCodeBlocks.java

@@ -38,6 +38,7 @@
 import org.springframework.data.util.Streamable;
 import org.springframework.javapoet.CodeBlock;
 import org.springframework.javapoet.CodeBlock.Builder;
+import org.springframework.util.ClassUtils;
 import org.springframework.util.NumberUtils;
 import org.springframework.util.StringUtils;
 
@@ -238,9 +239,25 @@ static void appendReadPreference(AotQueryMethodGenerationContext context, Builde
 	 * {@link MethodReturn} indicates so.
 	 */
 	public static CodeBlock potentiallyWrapStreamable(MethodReturn methodReturn, CodeBlock returningIterable) {
-		return methodReturn.toClass().equals(Streamable.class)
-				? CodeBlock.of("$T.of($L)", Streamable.class, returningIterable)
-				: returningIterable;
+
+		Class<?> returnType = methodReturn.toClass();
+		if (returnType.equals(Streamable.class)) {
+			return CodeBlock.of("$T.of($L)", Streamable.class, returningIterable);
+		}
+		if (ClassUtils.isAssignable(Streamable.class, returnType)) {
+			CodeBlock streamable = CodeBlock.of("$T.of($L)", Streamable.class, returningIterable);
+			if (ClassUtils.hasConstructor(returnType, Streamable.class)) {
+				return CodeBlock.of("new $T($L)", returnType, streamable);
+			}
+			if (ClassUtils.hasAtLeastOneMethodWithName(returnType, "of")) {
+				return CodeBlock.of("$T.of($L)", returnType, streamable);
+			}
+			if (ClassUtils.hasAtLeastOneMethodWithName(returnType, "valueOf")) {
+				return CodeBlock.of("$T.valueOf($L)", returnType, streamable);
+			}
+		}
+
+		return returningIterable;
 	}
 
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/AbstractPersonRepositoryIntegrationTests.java

@@ -73,6 +73,7 @@
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
 import org.springframework.data.mongodb.repository.Person.Sex;
+import org.springframework.data.mongodb.repository.PersonRepository.Persons;
 import org.springframework.data.mongodb.repository.SampleEvaluationContextExtension.SampleSecurityContextHolder;
 import org.springframework.data.mongodb.test.util.DirtiesStateExtension;
 import org.springframework.data.mongodb.test.util.DirtiesStateExtension.DirtiesState;
@@ -324,6 +325,17 @@ void streamPersonByAddressCorrectly() {
 		assertThat(result).hasSize(1).contains(dave);
 	}
 
+	@Test // GH-5089
+	void useCustomReturnTypeImplementingStreamable() {
+
+		Address address = new Address("Foo Street 1", "C0123", "Bar");
+		dave.setAddress(address);
+		repository.save(dave);
+
+		Persons result = repository.streamPersonsByAddress(address);
+		assertThat(result).hasSize(1).contains(dave);
+	}
+
 	@Test // GH-5089
 	void streamPersonByAddressCorrectlyWhenPaged() {
 

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/PersonRepository.java

@@ -17,6 +17,7 @@
 
 import java.util.Collection;
 import java.util.Date;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
@@ -214,6 +215,8 @@ Window<Person> findByLastnameLikeOrderByLastnameAscFirstnameAsc(Pattern lastname
 
 	Streamable<Person> streamByAddress(Address address);
 
+	Persons streamPersonsByAddress(Address address);
+
 	Streamable<Person> streamByAddress(Address address, Pageable pageable);
 
 	List<Person> findByAddressZipCode(String zipCode);
@@ -502,4 +505,17 @@ Person findPersonByManyArguments(String firstname, String lastname, String email
 
 	List<Person> findBySpiritAnimal(User user);
 
+	class Persons implements Streamable<Person> {
+
+		private final Streamable<Person> streamable;
+
+		public Persons(Streamable<Person> streamable) {
+			this.streamable = streamable;
+		}
+
+		@Override
+		public Iterator<Person> iterator() {
+			return streamable.iterator();
+		}
+	}
 }