JAVA-2949: Provide mapper support for CompletionStage<Stream<T>> (#1563)

Author: adutra

changelog/README.md

@@ -4,6 +4,7 @@
 
 ### 4.11.3 (in progress)
 
+- [bug] JAVA-2949: Provide mapper support for CompletionStage<Stream<T>>
 - [bug] JAVA-2950: Remove reference to Reflection class from DependencyCheck
 
 ### 4.11.2

integration-tests/src/test/java/com/datastax/oss/driver/mapper/QueryReturnTypesIT.java

@@ -47,6 +47,7 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.ExecutionException;
 import java.util.stream.Stream;
 import org.junit.Before;
 import org.junit.BeforeClass;
@@ -239,6 +240,13 @@ public void should_execute_async_query_and_map_to_iterable() {
     assertThat(iterable.hasMorePages()).isFalse();
   }
 
+  @Test
+  public void should_execute_query_and_map_to_stream_async()
+      throws ExecutionException, InterruptedException {
+    CompletableFuture<Stream<TestEntity>> stream = dao.findByIdAsStreamAsync(1);
+    assertThat(stream.get()).hasSize(10);
+  }
+
   @Dao
   @DefaultNullSavingStrategy(NullSavingStrategy.SET_TO_NULL)
   public interface TestDao {
@@ -300,6 +308,9 @@ public interface TestDao {
 
     @Query("SELECT * FROM ${qualifiedTableId} WHERE id = :id")
     CompletableFuture<MappedAsyncPagingIterable<TestEntity>> findByIdAsync(int id);
+
+    @Query("SELECT * FROM ${qualifiedTableId} WHERE id = :id")
+    CompletableFuture<Stream<TestEntity>> findByIdAsStreamAsync(int id);
   }
 
   @Entity

integration-tests/src/test/java/com/datastax/oss/driver/mapper/SelectIT.java

@@ -18,6 +18,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.MappedAsyncPagingIterable;
 import com.datastax.oss.driver.api.core.PagingIterable;
 import com.datastax.oss.driver.api.core.cql.SimpleStatement;
 import com.datastax.oss.driver.api.mapper.annotations.Dao;
@@ -100,11 +101,21 @@ public void should_select_all() {
     assertThat(dao.all().all()).hasSize(2);
   }
 
+  @Test
+  public void should_select_all_async() {
+    assertThat(CompletableFutures.getUninterruptibly(dao.allAsync()).currentPage()).hasSize(2);
+  }
+
   @Test
   public void should_select_all_stream() {
     assertThat(dao.stream()).hasSize(2);
   }
 
+  @Test
+  public void should_select_all_stream_async() {
+    assertThat(CompletableFutures.getUninterruptibly(dao.streamAsync())).hasSize(2);
+  }
+
   @Test
   public void should_select_by_primary_key_asynchronously() {
     assertThat(CompletableFutures.getUninterruptibly(dao.findByIdAsync(FLAMETHROWER.getId())))
@@ -211,9 +222,15 @@ public interface ProductDao {
     @Select
     PagingIterable<Product> all();
 
+    @Select
+    CompletionStage<MappedAsyncPagingIterable<Product>> allAsync();
+
     @Select
     Stream<Product> stream();
 
+    @Select
+    CompletionStage<Stream<Product>> streamAsync();
+
     @Select
     Optional<Product> findOptionalById(UUID productId);
 

manual/mapper/daos/select/README.md

@@ -135,6 +135,11 @@ In all cases, the method can return:
     @Select(customWhereClause = "description LIKE :searchString")
     CompletionStage<MappedAsyncPagingIterable<Product>> findByDescriptionAsync(String searchString);
     ```
+
+    For streams, even if the initial query is executed asynchronously, traversing the returned
+    stream may block the traversing thread. Blocking calls can indeed be required as more results
+    are fetched from the server in the background. For this reason, _the usage of
+    `CompletionStage<Stream<T>>` cannot be considered as a fully asynchronous execution method_.
   
 * a [MappedReactiveResultSet] of the entity class.
 

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoSelectMethodGenerator.java

@@ -20,6 +20,7 @@
 import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.FUTURE_OF_ASYNC_PAGING_ITERABLE;
 import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.FUTURE_OF_ENTITY;
 import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.FUTURE_OF_OPTIONAL_ENTITY;
+import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.FUTURE_OF_STREAM;
 import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.MAPPED_REACTIVE_RESULT_SET;
 import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.OPTIONAL_ENTITY;
 import static com.datastax.oss.driver.internal.mapper.processor.dao.DefaultDaoReturnTypeKind.PAGING_ITERABLE;
@@ -71,6 +72,7 @@ protected Set<DaoReturnTypeKind> getSupportedReturnTypes() {
         PAGING_ITERABLE,
         STREAM,
         FUTURE_OF_ASYNC_PAGING_ITERABLE,
+        FUTURE_OF_STREAM,
         MAPPED_REACTIVE_RESULT_SET,
         CUSTOM);
   }

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DefaultDaoReturnTypeKind.java

@@ -475,6 +475,26 @@ public CodeBlock wrapWithErrorHandling(
       return innerBlock;
     }
   },
+
+  FUTURE_OF_STREAM {
+    @Override
+    public void addExecuteStatement(
+        CodeBlock.Builder methodBuilder,
+        String helperFieldName,
+        ExecutableElement methodElement,
+        Map<Name, TypeElement> typeParameters) {
+      methodBuilder.addStatement(
+          "return executeAsyncAndMapToEntityStream(boundStatement, $L)", helperFieldName);
+    }
+
+    @Override
+    public CodeBlock wrapWithErrorHandling(
+        CodeBlock innerBlock,
+        ExecutableElement methodElement,
+        Map<Name, TypeElement> typeParameters) {
+      return wrapWithErrorHandling(innerBlock, FAILED_FUTURE);
+    }
+  },
   ;
 
   @Override

mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/dao/DefaultDaoReturnTypeParser.java

@@ -101,6 +101,7 @@ public class DefaultDaoReturnTypeParser implements DaoReturnTypeParser {
               .put(
                   MappedAsyncPagingIterable.class,
                   DefaultDaoReturnTypeKind.FUTURE_OF_ASYNC_PAGING_ITERABLE)
+              .put(Stream.class, DefaultDaoReturnTypeKind.FUTURE_OF_STREAM)
               .build();
 
   protected final ProcessorContext context;

mapper-processor/src/test/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoQueryMethodGeneratorTest.java

@@ -43,9 +43,11 @@ public static Object[][] invalidSignatures() {
       {
         "Invalid return type: Query methods must return one of [VOID, BOOLEAN, LONG, ROW, "
             + "ENTITY, OPTIONAL_ENTITY, RESULT_SET, BOUND_STATEMENT, PAGING_ITERABLE, FUTURE_OF_VOID, "
-            + "FUTURE_OF_BOOLEAN, FUTURE_OF_LONG, FUTURE_OF_ROW, FUTURE_OF_ENTITY, "
-            + "FUTURE_OF_OPTIONAL_ENTITY, FUTURE_OF_ASYNC_RESULT_SET, "
-            + "FUTURE_OF_ASYNC_PAGING_ITERABLE, REACTIVE_RESULT_SET, MAPPED_REACTIVE_RESULT_SET, STREAM]",
+            + "FUTURE_OF_BOOLEAN, FUTURE_OF_LONG, FUTURE_OF_ROW, "
+            + "FUTURE_OF_ENTITY, FUTURE_OF_OPTIONAL_ENTITY, "
+            + "FUTURE_OF_ASYNC_RESULT_SET, FUTURE_OF_ASYNC_PAGING_ITERABLE, "
+            + "REACTIVE_RESULT_SET, MAPPED_REACTIVE_RESULT_SET, "
+            + "STREAM, FUTURE_OF_STREAM]",
         MethodSpec.methodBuilder("select")
             .addAnnotation(
                 AnnotationSpec.builder(Query.class)

mapper-processor/src/test/java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoSelectMethodGeneratorTest.java

@@ -42,7 +42,7 @@ public static Object[][] invalidSignatures() {
       {
         "Invalid return type: Select methods must return one of [ENTITY, OPTIONAL_ENTITY, "
             + "FUTURE_OF_ENTITY, FUTURE_OF_OPTIONAL_ENTITY, PAGING_ITERABLE, STREAM, "
-            + "FUTURE_OF_ASYNC_PAGING_ITERABLE, MAPPED_REACTIVE_RESULT_SET]",
+            + "FUTURE_OF_ASYNC_PAGING_ITERABLE, FUTURE_OF_STREAM, MAPPED_REACTIVE_RESULT_SET]",
         MethodSpec.methodBuilder("select")
             .addAnnotation(Select.class)
             .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
@@ -52,7 +52,7 @@ public static Object[][] invalidSignatures() {
       {
         "Invalid return type: Select methods must return one of [ENTITY, OPTIONAL_ENTITY, "
             + "FUTURE_OF_ENTITY, FUTURE_OF_OPTIONAL_ENTITY, PAGING_ITERABLE, STREAM, "
-            + "FUTURE_OF_ASYNC_PAGING_ITERABLE, MAPPED_REACTIVE_RESULT_SET]",
+            + "FUTURE_OF_ASYNC_PAGING_ITERABLE, FUTURE_OF_STREAM, MAPPED_REACTIVE_RESULT_SET]",
         MethodSpec.methodBuilder("select")
             .addAnnotation(Select.class)
             .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)

mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DaoBase.java

@@ -35,6 +35,7 @@
 import com.datastax.oss.driver.api.mapper.entity.saving.NullSavingStrategy;
 import com.datastax.oss.driver.internal.core.ConsistencyLevelRegistry;
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
+import com.datastax.oss.driver.internal.core.cql.ResultSets;
 import com.datastax.oss.protocol.internal.ProtocolConstants;
 import java.time.Duration;
 import java.util.Optional;
@@ -290,6 +291,13 @@ CompletableFuture<MappedAsyncPagingIterable<EntityT>> executeAsyncAndMapToEntity
     return executeAsync(statement).thenApply(rs -> rs.map(entityHelper::get));
   }
 
+  protected <EntityT> CompletableFuture<Stream<EntityT>> executeAsyncAndMapToEntityStream(
+      Statement<?> statement, EntityHelper<EntityT> entityHelper) {
+    return executeAsync(statement)
+        .thenApply(ResultSets::newInstance)
+        .thenApply(rs -> StreamSupport.stream(rs.map(entityHelper::get).spliterator(), false));
+  }
+
   protected static void throwIfProtocolVersionV3(MapperContext context) {
     if (isProtocolVersionV3(context)) {
       throw new MapperException(