DATACASS-809 - Add queryForStream methods to CqlOperations.

CqlOperations now exposes queryForStream methods to consume a query as java.util.stream.Stream accepting a RowMapper.

Author: mp911de

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraTemplate.java

@@ -19,7 +19,6 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.Stream;
-import java.util.stream.StreamSupport;
 
 import org.springframework.beans.BeansException;
 import org.springframework.context.ApplicationContext;
@@ -374,10 +373,8 @@ public <T> Stream<T> stream(Statement<?> statement, Class<T> entityClass) throws
 		Assert.notNull(statement, "Statement must not be null");
 		Assert.notNull(entityClass, "Entity type must not be null");
 
-		ResultSet resultSet = getCqlOperations().queryForResultSet(statement);
-
-		return StreamSupport.stream(resultSet.spliterator(), false)
-				.map(getMapper(entityClass, entityClass, EntityQueryUtils.getTableName(statement)));
+		Function<Row, T> mapper = getMapper(entityClass, entityClass, EntityQueryUtils.getTableName(statement));
+		return getCqlOperations().queryForStream(statement, (row, rowNum) -> mapper.apply(row));
 	}
 
 	// -------------------------------------------------------------------------
@@ -453,10 +450,8 @@ <T> Stream<T> doStream(Query query, Class<?> entityClass, CqlIdentifier tableNam
 		StatementBuilder<Select> select = getStatementFactory().select(query, getRequiredPersistentEntity(entityClass),
 				tableName);
 
-		ResultSet resultSet = getCqlOperations().queryForResultSet(select.build());
-
 		Function<Row, T> mapper = getMapper(entityClass, returnType, tableName);
-		return StreamSupport.stream(resultSet.map(mapper).spliterator(), false);
+		return getCqlOperations().queryForStream(select.build(), (row, rowNum) -> mapper.apply(row));
 	}
 
 	/* (non-Javadoc)

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/cql/AsyncCqlTemplate.java

@@ -752,8 +752,7 @@ protected AsyncPreparedStatementCreator newAsyncPreparedStatementCreator(String
 	}
 
 	/**
-	 * Constructs a new instance of the {@link ResultSetExtractor} initialized with and adapting the given
-	 * {@link RowCallbackHandler}.
+	 * Constructs a new instance of the {@link ResultSetExtractor} adapting the given {@link RowCallbackHandler}.
 	 *
 	 * @param rowCallbackHandler {@link RowCallbackHandler} to adapt as a {@link ResultSetExtractor}.
 	 * @return a {@link ResultSetExtractor} implementation adapting an instance of the {@link RowCallbackHandler}.
@@ -767,8 +766,7 @@ protected AsyncRowCallbackHandlerResultSetExtractor newAsyncResultSetExtractor(
 	}
 
 	/**
-	 * Constructs a new instance of the {@link ResultSetExtractor} initialized with and adapting the given
-	 * {@link RowMapper}.
+	 * Constructs a new instance of the {@link ResultSetExtractor} adapting the given {@link RowMapper}.
 	 *
 	 * @param rowMapper {@link RowMapper} to adapt as a {@link ResultSetExtractor}.
 	 * @return a {@link ResultSetExtractor} implementation adapting an instance of the {@link RowMapper}.

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/cql/CqlOperations.java

@@ -18,16 +18,17 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Stream;
+
+import org.springframework.dao.DataAccessException;
+import org.springframework.dao.IncorrectResultSizeDataAccessException;
+import org.springframework.lang.Nullable;
 
 import com.datastax.oss.driver.api.core.cql.PreparedStatement;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
 import com.datastax.oss.driver.api.core.cql.Row;
 import com.datastax.oss.driver.api.core.cql.Statement;
 
-import org.springframework.dao.DataAccessException;
-import org.springframework.dao.IncorrectResultSizeDataAccessException;
-import org.springframework.lang.Nullable;
-
 /**
  * Interface specifying a basic set of CQL operations. Implemented by {@link CqlTemplate}. Not often used directly, but
  * a useful option to enhance testability, as it can easily be mocked or stubbed.
@@ -195,6 +196,20 @@ public interface CqlOperations {
 	 */
 	<T> List<T> query(String cql, RowMapper<T> rowMapper, Object... args) throws DataAccessException;
 
+	/**
+	 * Query given CQL to create a prepared statement from CQL and a list of arguments to bind to the query, mapping each
+	 * row to a Java object via a {@link RowMapper} and turning it into an iterable {@link Stream}.
+	 *
+	 * @param cql static CQL to execute, must not be empty or {@literal null}.
+	 * @param rowMapper object that will map one object per row
+	 * @param args arguments to bind to the query (leaving it to the {@link PreparedStatement} to guess the corresponding
+	 *          CQL type)
+	 * @return the result {@link Stream}, containing mapped objects
+	 * @throws DataAccessException if there is any problem executing the query.
+	 * @since 3.1
+	 */
+	<T> Stream<T> queryForStream(String cql, RowMapper<T> rowMapper, Object... args) throws DataAccessException;
+
 	/**
 	 * Query using a prepared statement, reading the {@link ResultSet} with a {@link ResultSetExtractor}.
 	 *
@@ -537,6 +552,22 @@ <T> List<T> query(String cql, @Nullable PreparedStatementBinder psb, RowMapper<T
 	 */
 	<T> List<T> query(Statement<?> statement, RowMapper<T> rowMapper) throws DataAccessException;
 
+	/**
+	 * Execute a query given static CQL, mapping each row to a Java object via a {@link RowMapper} and turning it into an
+	 * iterable {@link Stream}.
+	 * <p>
+	 * Uses a CQL Statement, not a {@link PreparedStatement}. If you want to execute a static query with a
+	 * {@link PreparedStatement}, use the overloaded {@code query} method with {@literal null} as argument array.
+	 *
+	 * @param statement static CQL {@link Statement}, must not be {@literal null}.
+	 * @param rowMapper object that will map one object per row, must not be {@literal null}.
+	 * @return the result {@link Stream}, containing mapped objects.
+	 * @throws DataAccessException if there is any problem executing the query.
+	 * @since 3.1
+	 * @see #queryForStream(String, RowMapper, Object...)
+	 */
+	<T> Stream<T> queryForStream(Statement<?> statement, RowMapper<T> rowMapper) throws DataAccessException;
+
 	/**
 	 * Execute a query for a result {@link List}, given static CQL.
 	 * <p>
@@ -729,6 +760,20 @@ void query(PreparedStatementCreator preparedStatementCreator, RowCallbackHandler
 	<T> List<T> query(PreparedStatementCreator preparedStatementCreator, RowMapper<T> rowMapper)
 			throws DataAccessException;
 
+	/**
+	 * Query using a prepared statement, mapping each row to a Java object via a {@link RowMapper} and turning it into an
+	 * iterable {@link Stream}.
+	 *
+	 * @param preparedStatementCreator object that can create a {@link PreparedStatement} given a
+	 *          {@link com.datastax.oss.driver.api.core.CqlSession}, must not be {@literal null}.
+	 * @param rowMapper object that will map one object per row, must not be {@literal null}.
+	 * @return the result {@link Stream}, containing mapped objects.
+	 * @throws DataAccessException if there is any problem executing the query.
+	 * @since 3.1
+	 */
+	<T> Stream<T> queryForStream(PreparedStatementCreator preparedStatementCreator, RowMapper<T> rowMapper)
+			throws DataAccessException;
+
 	/**
 	 * Query using a prepared statement and a {@link PreparedStatementBinder} implementation that knows how to bind values
 	 * to the query, reading the {@link ResultSet} with a {@link ResultSetExtractor}.
@@ -777,6 +822,24 @@ void query(PreparedStatementCreator preparedStatementCreator, @Nullable Prepared
 	<T> List<T> query(PreparedStatementCreator preparedStatementCreator, @Nullable PreparedStatementBinder psb,
 			RowMapper<T> rowMapper) throws DataAccessException;
 
+	/**
+	 * Query using a prepared statement and a {@link PreparedStatementBinder} implementation that knows how to bind values
+	 * to the query, mapping each row to a Java object via a {@link RowMapper} and turning it into an iterable
+	 * {@link Stream}.
+	 *
+	 * @param preparedStatementCreator object that can create a {@link PreparedStatement} given a
+	 *          {@link com.datastax.oss.driver.api.core.CqlSession}, must not be {@literal null}.
+	 * @param psb object that knows how to set values on the prepared statement. If this is {@literal null}, the CQL will
+	 *          be assumed to contain no bind parameters. Even if there are no bind parameters, this object may be used to
+	 *          set fetch size and other performance options.
+	 * @param rowMapper object that will map one object per row, must not be {@literal null}.
+	 * @return the result {@link Stream}, containing mapped objects.
+	 * @throws DataAccessException if there is any problem executing the query.
+	 * @since 3.1
+	 */
+	<T> Stream<T> queryForStream(PreparedStatementCreator preparedStatementCreator, @Nullable PreparedStatementBinder psb,
+			RowMapper<T> rowMapper) throws DataAccessException;
+
 	// -------------------------------------------------------------------------
 	// Methods dealing with cluster metadata
 	// -------------------------------------------------------------------------

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/cql/CqlTemplate.java

@@ -18,7 +18,11 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Spliterator;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
 import org.springframework.dao.DataAccessException;
@@ -316,6 +320,16 @@ public <T> List<T> query(Statement<?> statement, RowMapper<T> rowMapper) throws
 		return query(statement, newResultSetExtractor(rowMapper));
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.cassandra.core.cqlOperations#queryForStream(com.datastax.oss.driver.api.core.cql.Statement, org.springframework.data.cassandra.core.cql.RowMapper)
+	 */
+	@Override
+	public <T> Stream<T> queryForStream(Statement<?> statement, RowMapper<T> rowMapper) throws DataAccessException {
+		// noinspection ConstantConditions
+		return query(statement, newStreamExtractor(rowMapper));
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.cassandra.core.cqlOperations#queryForList(com.datastax.oss.driver.api.core.cql.Statement)
@@ -342,7 +356,6 @@ public <T> List<T> queryForList(Statement<?> statement, Class<T> elementType) th
 	 */
 	@Override
 	public Map<String, Object> queryForMap(Statement<?> statement) throws DataAccessException {
-		// noinspection ConstantConditions
 		return queryForObject(statement, newColumnMapRowMapper());
 	}
 
@@ -485,6 +498,17 @@ public <T> List<T> query(PreparedStatementCreator preparedStatementCreator, RowM
 		return query(preparedStatementCreator, null, newResultSetExtractor(rowMapper));
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.cassandra.core.cqlOperations#query(org.springframework.data.cassandra.core.cql.PreparedStatementCreator, org.springframework.data.cassandra.core.cql.RowMapper)
+	 */
+	@Override
+	public <T> Stream<T> queryForStream(PreparedStatementCreator preparedStatementCreator, RowMapper<T> rowMapper)
+			throws DataAccessException {
+		// noinspection ConstantConditions
+		return query(preparedStatementCreator, null, newStreamExtractor(rowMapper));
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.cassandra.core.cqlOperations#query(org.springframework.data.cassandra.core.cql.PreparedStatementCreator, org.springframework.data.cassandra.core.cql.PreparedStatementBinder, org.springframework.data.cassandra.core.cql.ResultSetExtractor)
@@ -544,6 +568,17 @@ public <T> List<T> query(PreparedStatementCreator preparedStatementCreator, @Nul
 		return query(preparedStatementCreator, psb, newResultSetExtractor(rowMapper));
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.cassandra.core.cqlOperations#queryForStream(org.springframework.data.cassandra.core.cql.PreparedStatementCreator, org.springframework.data.cassandra.core.cql.PreparedStatementBinder, org.springframework.data.cassandra.core.cql.RowMapper)
+	 */
+	@Override
+	public <T> Stream<T> queryForStream(PreparedStatementCreator preparedStatementCreator,
+			@Nullable PreparedStatementBinder psb, RowMapper<T> rowMapper) throws DataAccessException {
+		// noinspection ConstantConditions
+		return query(preparedStatementCreator, psb, newStreamExtractor(rowMapper));
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.cassandra.core.cqlOperations#query(java.lang.String, org.springframework.data.cassandra.core.cql.ResultSetExtractor, java.lang.Object[])
@@ -574,6 +609,16 @@ public <T> List<T> query(String cql, RowMapper<T> rowMapper, Object... args) thr
 		return query(newPreparedStatementCreator(cql), newPreparedStatementBinder(args), newResultSetExtractor(rowMapper));
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.cassandra.core.cqlOperations#queryForStream(java.lang.String, org.springframework.data.cassandra.core.cql.RowMapper, java.lang.Object[])
+	 */
+	@Override
+	public <T> Stream<T> queryForStream(String cql, RowMapper<T> rowMapper, Object... args) throws DataAccessException {
+		// noinspection ConstantConditions
+		return query(newPreparedStatementCreator(cql), newPreparedStatementBinder(args), newStreamExtractor(rowMapper));
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.cassandra.core.cqlOperations#query(java.lang.String, org.springframework.data.cassandra.core.cql.PreparedStatementBinder, org.springframework.data.cassandra.core.cql.ResultSetExtractor)
@@ -732,8 +777,7 @@ protected PreparedStatementCreator newPreparedStatementCreator(String cql) {
 	}
 
 	/**
-	 * Constructs a new instance of the {@link ResultSetExtractor} initialized with and adapting the given
-	 * {@link RowCallbackHandler}.
+	 * Constructs a new instance of the {@link ResultSetExtractor} adapting the given {@link RowCallbackHandler}.
 	 *
 	 * @param rowCallbackHandler {@link RowCallbackHandler} to adapt as a {@link ResultSetExtractor}.
 	 * @return a {@link ResultSetExtractor} implementation adapting an instance of the {@link RowCallbackHandler}.
@@ -746,8 +790,7 @@ protected RowCallbackHandlerResultSetExtractor newResultSetExtractor(RowCallback
 	}
 
 	/**
-	 * Constructs a new instance of the {@link ResultSetExtractor} initialized with and adapting the given
-	 * {@link RowMapper}.
+	 * Constructs a new instance of the {@link ResultSetExtractor} adapting the given {@link RowMapper}.
 	 *
 	 * @param rowMapper {@link RowMapper} to adapt as a {@link ResultSetExtractor}.
 	 * @return a {@link ResultSetExtractor} implementation adapting an instance of the {@link RowMapper}.
@@ -760,8 +803,7 @@ protected <T> RowMapperResultSetExtractor<T> newResultSetExtractor(RowMapper<T>
 	}
 
 	/**
-	 * Constructs a new instance of the {@link ResultSetExtractor} initialized with and adapting the given
-	 * {@link RowMapper}.
+	 * Constructs a new instance of the {@link ResultSetExtractor} adapting the given {@link RowMapper}.
 	 *
 	 * @param rowMapper {@link RowMapper} to adapt as a {@link ResultSetExtractor}.
 	 * @param rowsExpected number of expected rows in the {@link ResultSet}.
@@ -774,6 +816,19 @@ protected <T> RowMapperResultSetExtractor<T> newResultSetExtractor(RowMapper<T>
 		return new RowMapperResultSetExtractor<>(rowMapper, rowsExpected);
 	}
 
+	/**
+	 * Constructs a new instance of the {@link ResultSetExtractor} adapting the given {@link RowMapper}.
+	 *
+	 * @param rowMapper {@link RowMapper} to adapt as a {@link ResultSetExtractor}.
+	 * @return a {@link ResultSetExtractor} implementation adapting an instance of the {@link RowMapper}.
+	 * @see ResultSetExtractor
+	 * @see RowCallbackHandler
+	 * @since 3.1
+	 */
+	protected <T> ResultSetExtractor<Stream<T>> newStreamExtractor(RowMapper<T> rowMapper) {
+		return resultSet -> new ResultSetSpliterator<>(resultSet, rowMapper).stream();
+	}
+
 	private CqlSession getCurrentSession() {
 
 		SessionFactory sessionFactory = getSessionFactory();
@@ -801,10 +856,79 @@ protected RowCallbackHandlerResultSetExtractor(RowCallbackHandler rowCallbackHan
 		@Nullable
 		public Object extractData(ResultSet resultSet) {
 
-			StreamSupport.stream(resultSet.spliterator(), false).forEach(rowCallbackHandler::processRow);
-
+			resultSet.forEach(rowCallbackHandler::processRow);
 			return null;
 		}
 	}
 
+	/**
+	 * Spliterator for queryForStream adaptation of a {@link ResultSet} to a {@link Stream}.
+	 *
+	 * @since 3.1
+	 */
+	private static class ResultSetSpliterator<T> implements Spliterator<T> {
+
+		private final Spliterator<Row> delegate;
+
+		private final RowMapper<T> rowMapper;
+
+		private final AtomicInteger counter;
+
+		public ResultSetSpliterator(ResultSet rs, RowMapper<T> rowMapper) {
+			this.delegate = rs.spliterator();
+			this.rowMapper = rowMapper;
+			this.counter = new AtomicInteger();
+		}
+
+		private ResultSetSpliterator(Spliterator<Row> delegate, RowMapper<T> rowMapper, AtomicInteger counter) {
+			this.delegate = delegate;
+			this.rowMapper = rowMapper;
+			this.counter = counter;
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see java.util.Spliterator#tryAdvance(java.util.function.Consumer)
+		 */
+		@Override
+		public boolean tryAdvance(Consumer<? super T> action) {
+			return this.delegate.tryAdvance(row -> action.accept(this.rowMapper.mapRow(row, this.counter.incrementAndGet())));
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see java.util.Spliterator#trySplit()
+		 */
+		@Override
+		@Nullable
+		public Spliterator<T> trySplit() {
+			return new ResultSetSpliterator<>(delegate.trySplit(), this.rowMapper, this.counter);
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see java.util.Spliterator#estimateSize()
+		 */
+		@Override
+		public long estimateSize() {
+			return Long.MAX_VALUE;
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see java.util.Spliterator#characteristics()
+		 */
+		@Override
+		public int characteristics() {
+			return Spliterator.ORDERED;
+		}
+
+		/**
+		 * @return
+		 */
+		public Stream<T> stream() {
+			return StreamSupport.stream(this, false);
+		}
+	}
+
 }