Add AOT support for Streamable wrapper support.

We now support streamable wrappers through AOT repositories.

See #1620

Author: mp911de

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/aot/CassandraCodeBlocks.java

@@ -23,6 +23,8 @@
 import org.jspecify.annotations.Nullable;
 
 import org.springframework.core.annotation.MergedAnnotation;
+import org.springframework.core.convert.TypeDescriptor;
+import org.springframework.core.convert.support.DefaultConversionService;
 import org.springframework.data.cassandra.core.CassandraOperations;
 import org.springframework.data.cassandra.core.ExecutableSelectOperation;
 import org.springframework.data.cassandra.core.cql.QueryOptions;
@@ -718,13 +720,29 @@ CodeBlock build() {
 
 				if (streamableResult && Collection.class.isAssignableFrom(methodReturn.toClass())) {
 					result = CodeBlock.of("$L.getContent()", result);
-				} else if (!streamableResult && methodReturn.toClass().equals(Streamable.class)) {
+				} else if (!streamableResult && (isStreamable(methodReturn) || isStreamableWrapper(methodReturn))) {
 					result = CodeBlock.of("$T.of(($T) $L)", Streamable.class, Iterable.class, result);
 				}
 
-				builder.addStatement(returnBuilder //
+				if (isStreamableWrapper(methodReturn) && canConvert(Streamable.class, methodReturn)) {
+
+					Builder wrapperBuilder = CodeBlock.builder();
+
+					builder.addStatement("$1T<$2T> $3L = $4L", Streamable.class, actualReturnType,
+							context.localVariable("streamable"), result);
+
+					builder.addStatement(
+							"return ($1T) $2T.getSharedInstance().convert($3L, $4T.valueOf($5T.class), $4T.valueOf($6T.class))",
+							methodReturn.getTypeName(), DefaultConversionService.class, context.localVariable("streamable"),
+							TypeDescriptor.class, Streamable.class, methodReturn.toClass());
+
+					builder.add(wrapperBuilder.build());
+				} else {
+
+					builder.addStatement(returnBuilder //
 						.optional("($T) $L", methodReturn.getTypeName(), result) //
 						.build());
+				}
 			} else {
 
 				CodeBlock executionBlock = execution.build();
@@ -733,16 +751,40 @@ CodeBlock build() {
 					returnBuilder.whenPrimitiveOrBoxed(Integer.class, "(int) $L", executionBlock);
 				} else if (streamableResult && Collection.class.isAssignableFrom(methodReturn.toClass())) {
 					executionBlock = CodeBlock.of("$L.getContent()", executionBlock);
-				} else if (!streamableResult && methodReturn.toClass().equals(Streamable.class)) {
+				} else if (!streamableResult && (isStreamable(methodReturn) || isStreamableWrapper(methodReturn))) {
 					executionBlock = CodeBlock.of("$T.of($L)", Streamable.class, executionBlock);
 				}
 
-				builder.addStatement(returnBuilder.optional(executionBlock) //
+				if (isStreamableWrapper(methodReturn) && canConvert(Streamable.class, methodReturn)) {
+
+					builder.addStatement("$1T<$2T> $3L = $4L", Streamable.class, actualReturnType,
+							context.localVariable("streamable"), executionBlock);
+
+					builder.addStatement(
+							"return ($1T) $2T.getSharedInstance().convert($3L, $4T.valueOf($5T.class), $4T.valueOf($6T.class))",
+							methodReturn.getTypeName(), DefaultConversionService.class, context.localVariable("streamable"),
+							TypeDescriptor.class, Streamable.class, methodReturn.toClass());
+				} else {
+
+					builder.addStatement(returnBuilder.optional(executionBlock) //
 						.build());
+				}
 			}
 
 			return builder.build();
 		}
+
+		private boolean canConvert(Class<?> from, MethodReturn methodReturn) {
+			return DefaultConversionService.getSharedInstance().canConvert(from, methodReturn.toClass());
+		}
+
+		private static boolean isStreamable(MethodReturn methodReturn) {
+			return methodReturn.toClass().equals(Streamable.class);
+		}
+
+		private static boolean isStreamableWrapper(MethodReturn methodReturn) {
+			return !isStreamable(methodReturn) && Streamable.class.isAssignableFrom(methodReturn.toClass());
+		}
 	}
 
 }

spring-data-cassandra/src/test/java/org/springframework/data/cassandra/repository/aot/CassandraRepositoryContributorIntegrationTests.java

@@ -216,6 +216,24 @@ void shouldConvertResultToStreamableWhenPageableParameterIsUsed() {
 			.containsExactly("Flynn (Walter Jr.)", "Skyler");
 	}
 
+	@Test // GH-1620
+	void shouldConvertResultToStreamableWrapper() {
+
+		assertThat(fragment.findWrapperByLastname("White")) //
+				.isInstanceOf(PersonRepository.People.class) //
+				.extracting(Person::getFirstname) //
+				.hasSize(3);
+	}
+
+	@Test // GH-1620
+	void shouldConvertResultToStreamableWrapperWhenPageableParameterIsUsed() {
+
+		assertThat(fragment.findWrapperByLastname("White", PageRequest.of(0, 2, Sort.by("firstname"))))
+				.isInstanceOf(PersonRepository.People.class) //
+				.extracting(Person::getFirstname) //
+				.containsExactly("Flynn (Walter Jr.)", "Skyler");
+	}
+
 	@Test // GH-1566
 	void shouldFindByFirstnameContains() {
 
@@ -331,6 +349,15 @@ void shouldQueryDeclaredWindow() {
 		assertThat(second.hasNext()).isFalse();
 	}
 
+	@Test // GH-1620
+	void shouldQueryDeclaredStreamable() {
+
+		PersonRepository.People people = fragment.findDeclaredWrapperByLastname("White",
+				PageRequest.of(0, 2, Sort.by("firstname")));
+
+		assertThat(people).hasSize(2);
+	}
+
 	@Test // GH-1566
 	void shouldFindNamedQuery() {
 

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

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.cassandra.repository.aot;
 
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -38,10 +39,10 @@
 import org.springframework.data.domain.Vector;
 import org.springframework.data.domain.Window;
 import org.springframework.data.repository.CrudRepository;
+import org.springframework.data.util.Streamable;
 
 import com.datastax.oss.driver.api.core.DefaultConsistencyLevel;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
-import org.springframework.data.util.Streamable;
 
 /**
  * AOT repository interface for {@link Person} entities.
@@ -68,6 +69,10 @@ public interface PersonRepository extends CrudRepository<Person, String> {
 
 	Streamable<Person> streamByLastname(String lastname, Pageable pageable);
 
+	People findWrapperByLastname(String lastname);
+
+	People findWrapperByLastname(String lastname, Pageable pageable);
+
 	List<Person> findByLastnameOrderByFirstnameAsc(String lastname);
 
 	Person findByFirstnameStartsWith(String prefix);
@@ -122,6 +127,9 @@ public interface PersonRepository extends CrudRepository<Person, String> {
 	Window<Person> findDeclaredWindowByLastname(String lastname, ScrollPosition scrollPosition, int sliceLimit,
 			Limit pageSize);
 
+	@Query(value = "select * from person where lastname = ?0 LIMIT 3")
+	People findDeclaredWrapperByLastname(String lastname, Pageable pageable);
+
 	// -------------------------------------------------------------------------
 	// Value Expressions
 	// -------------------------------------------------------------------------
@@ -240,4 +248,18 @@ public void setLastname(String lastname) {
 			this.lastname = lastname;
 		}
 	}
+
+	class People implements Streamable<Person> {
+
+		private final Streamable<Person> people;
+
+		public People(Streamable<Person> people) {
+			this.people = people;
+		}
+
+		@Override
+		public Iterator<Person> iterator() {
+			return people.iterator();
+		}
+	}
 }