DATACASS-297 - Add support for streaming queries to select entities.

We now support java.util.stream.Stream as return type on query methods. Streaming queries wrap the iterator returned by Cassandra's ResultSet. Elements are fetched in batches (see QueryOptions that can be set in CassandraCqlClusterFactoryBean) and processed element by element instead of reading all results. Streaming queries are supported only for select queries.

Original pull request: #65.

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

@@ -52,6 +52,7 @@ import org.springframework.dao.DataAccessException;
 import org.springframework.dao.IncorrectResultSizeDataAccessException;
 import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.dao.QueryTimeoutException;
+import org.springframework.dao.support.PersistenceExceptionTranslator;
 import org.springframework.util.Assert;
 
 import com.datastax.driver.core.BoundStatement;
@@ -854,19 +855,49 @@ public class CqlTemplate extends CassandraAccessor implements CqlOperations {
 	}
 
 	/**
-	 * Attempts to translate the {@link RuntimeException} into a Spring Data {@link Exception}.
+	 * Attempts to translate the {@link Exception} into a Spring Data {@link Exception}.
+	 * @param ex the Exception
+	 * @return the translated {@link RuntimeException}
 	 */
 	@SuppressWarnings("all")
-	protected RuntimeException translateExceptionIfPossible(RuntimeException e) {
-
-		RuntimeException resolved = getExceptionTranslator().translateExceptionIfPossible(e);
-		return (resolved != null ? resolved : e);
+	protected RuntimeException translateExceptionIfPossible(Exception ex) {
+		return translateExceptionIfPossible(ex, getExceptionTranslator());
 	}
 
+	/**
+	 * Tries to convert the given {@link RuntimeException} into a {@link DataAccessException} but returns the original
+	 * exception if the conversation failed. Thus allows safe re-throwing of the return value.
+	 *
+	 * @param ex the exception to translate
+	 * @param exceptionTranslator the {@link PersistenceExceptionTranslator} to be used for translation
+	 * @return
+	 */
 	@SuppressWarnings("all")
-	protected RuntimeException translateExceptionIfPossible(Exception e) {
-		return (e instanceof RuntimeException ? translateExceptionIfPossible((RuntimeException) e)
-				: new CassandraUncategorizedDataAccessException("Caught Uncategorized Exception", e));
+	protected static RuntimeException translateExceptionIfPossible(Exception ex, PersistenceExceptionTranslator exceptionTranslator) {
+
+		Assert.notNull(ex, "Exception must not be null");
+		Assert.notNull(exceptionTranslator, "PersistenceExceptionTranslator must not be null");
+
+		if (ex instanceof RuntimeException) {
+			return potentiallyConvertRuntimeException((RuntimeException) ex, exceptionTranslator);
+		}
+
+		return new CassandraUncategorizedDataAccessException("Caught Uncategorized Exception", ex);
+	}
+
+	/**
+	 * Tries to convert the given {@link RuntimeException} into a {@link DataAccessException} but returns the original
+	 * exception if the conversation failed. Thus allows safe re-throwing of the return value.
+	 *
+	 * @param ex the exception to translate
+	 * @param exceptionTranslator the {@link PersistenceExceptionTranslator} to be used for translation
+	 * @return
+	 */
+	private static RuntimeException potentiallyConvertRuntimeException(RuntimeException ex,
+			PersistenceExceptionTranslator exceptionTranslator) {
+
+		RuntimeException resolved = exceptionTranslator.translateExceptionIfPossible(ex);
+		return resolved == null ? ex : resolved;
 	}
 
 	@Override

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

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.cassandra.core;
 
+import java.util.Iterator;
 import java.util.List;
 
 import org.springframework.cassandra.core.Cancellable;
@@ -51,6 +52,20 @@ public interface CassandraOperations extends CqlOperations {
 	 */
 	CqlIdentifier getTableName(Class<?> entityClass);
 
+	/**
+	 * Executes the given select {@code query} on the entity table of the specified {@code type} backed by a Cassandra
+	 * {@link com.datastax.driver.core.ResultSet}.
+	 * <p>
+	 * Returns a {@link java.util.Iterator} that wraps the a Cassandra {@link com.datastax.driver.core.ResultSet}.
+	 * 
+	 * @param <T> element return type
+	 * @param query must not be empty and not {@literal null}.
+	 * @param type must not be {@literal null}.
+	 * @return
+	 * @since 1.5
+	 */
+	<T> Iterator<T> stream(String query, Class<T> type);
+
 	/**
 	 * Execute query and convert ResultSet to the list of entities.
 	 *

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

@@ -17,6 +17,7 @@ package org.springframework.data.cassandra.core;
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -34,6 +35,7 @@ import org.springframework.cassandra.core.util.CollectionUtils;
 import org.springframework.dao.DataAccessException;
 import org.springframework.dao.DuplicateKeyException;
 import org.springframework.dao.InvalidDataAccessApiUsageException;
+import org.springframework.dao.support.PersistenceExceptionTranslator;
 import org.springframework.data.cassandra.convert.CassandraConverter;
 import org.springframework.data.cassandra.convert.MappingCassandraConverter;
 import org.springframework.data.cassandra.mapping.CassandraMappingContext;
@@ -588,6 +590,29 @@ public class CassandraTemplate extends CqlTemplate implements CassandraOperation
 		return result;
 	}
 
+	/* (non-Javadoc)
+	 * @see org.springframework.data.cassandra.core.CassandraOperations#stream(java.lang.String, java.lang.Class)
+	 */
+	public <T> Iterator<T> stream(final String query, Class<T> type) {
+
+		Assert.hasText(query, "Query must not be empty");
+		Assert.notNull(type, "Type must not be null");
+
+		ResultSet resultSet = doExecute(new SessionCallback<ResultSet>() {
+
+			@Override
+			public ResultSet doInSession(Session s) throws DataAccessException {
+				return s.execute(query);
+			}
+		});
+
+		if (resultSet == null) {
+			return Collections.<T>emptyList().iterator();
+		}
+
+		return new ResultSetIteratorAdapter(resultSet.iterator(), getExceptionTranslator(), new CassandraConverterRowCallback<T>(cassandraConverter, type));
+	}
+
 	protected <T> List<T> select(final Select query, CassandraConverterRowCallback<T> readRowCallback) {
 
 		ResultSet resultSet = doExecute(new SessionCallback<ResultSet>() {
@@ -1067,8 +1092,7 @@ public class CassandraTemplate extends CqlTemplate implements CassandraOperation
 							throw new DuplicateKeyException("found two or more results in query " + query);
 						}
 						listener.onQueryComplete(result);
-					}
-					else{
+					} else {
 						listener.onQueryComplete(null);
 					}
 				} catch (Exception e) {
@@ -1085,4 +1109,38 @@ public class CassandraTemplate extends CqlTemplate implements CassandraOperation
 		throw new IllegalArgumentException(
 				String.format("Expected type String or Select; got type [%s] with value [%s]", query.getClass(), query));
 	}
+	
+	private static class ResultSetIteratorAdapter<T> implements Iterator<T>{
+
+		private final Iterator<Row> iterator;
+		private final PersistenceExceptionTranslator exceptionTranslator;
+		private final CassandraConverterRowCallback<T> rowCallback;
+
+		public ResultSetIteratorAdapter(Iterator<Row> iterator, PersistenceExceptionTranslator exceptionTranslator, CassandraConverterRowCallback<T> rowCallback) {
+			
+			this.iterator = iterator;
+			this.exceptionTranslator = exceptionTranslator;
+			this.rowCallback = rowCallback;
+		}
+
+		@Override
+		public boolean hasNext() {
+			
+			try {
+				return iterator.hasNext();
+			} catch (Exception e) {
+				throw translateExceptionIfPossible(e, exceptionTranslator);
+			}
+		}
+
+		@Override
+		public T next() {
+			
+			try {
+				return rowCallback.doWith(iterator.next());
+			} catch (Exception e) {
+				throw translateExceptionIfPossible(e, exceptionTranslator);
+			}
+		}
+	}
 }

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/AbstractCassandraQuery.java

@@ -34,6 +34,7 @@ import org.springframework.data.cassandra.repository.query.CassandraQueryExecuti
 import org.springframework.data.cassandra.repository.query.CassandraQueryExecution.ResultProcessingExecution;
 import org.springframework.data.cassandra.repository.query.CassandraQueryExecution.ResultSetQuery;
 import org.springframework.data.cassandra.repository.query.CassandraQueryExecution.SingleEntityExecution;
+import org.springframework.data.cassandra.repository.query.CassandraQueryExecution.StreamExecution;
 import org.springframework.data.repository.query.ParameterAccessor;
 import org.springframework.data.repository.query.RepositoryQuery;
 import org.springframework.data.repository.query.ResultProcessor;
@@ -100,13 +101,15 @@ public abstract class AbstractCassandraQuery implements RepositoryQuery {
 	private CassandraQueryExecution getExecution(String query, CassandraParameterAccessor accessor,
 			Converter<Object, Object> resultProcessing) {
 
-		return new ResultProcessingExecution(getExecutionToWrap(accessor), resultProcessing);
+		return new ResultProcessingExecution(getExecutionToWrap(accessor, resultProcessing), resultProcessing);
 	}
 
-	private CassandraQueryExecution getExecutionToWrap(CassandraParameterAccessor accessor) {
+	private CassandraQueryExecution getExecutionToWrap(CassandraParameterAccessor accessor, Converter<Object, Object> resultProcessing) {
 
 		if (method.isResultSetQuery()) {
 			return new ResultSetQuery(template);
+		} else if (method.isStreamQuery()) {
+			return new StreamExecution(template, resultProcessing);
 		} else if (method.isCollectionQuery()) {
 			return new CollectionExecution(template);
 		} else {

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/CassandraQueryExecution.java

@@ -16,10 +16,13 @@
 
 package org.springframework.data.cassandra.repository.query;
 
+import java.util.function.Function;
+
 import org.springframework.core.convert.converter.Converter;
 import org.springframework.data.cassandra.core.CassandraOperations;
 import org.springframework.data.repository.query.ResultProcessor;
 import org.springframework.data.repository.query.ReturnedType;
+import org.springframework.data.util.StreamUtils;
 import org.springframework.util.ClassUtils;
 
 import lombok.NonNull;
@@ -35,6 +38,33 @@ interface CassandraQueryExecution {
 
 	Object execute(String query, Class<?> type);
 
+	/**
+	 * {@link CassandraQueryExecution} for a Stream.
+	 *
+	 * @author Mark Paluch
+	 */
+	@RequiredArgsConstructor
+	final class StreamExecution implements CassandraQueryExecution {
+
+		private final @NonNull CassandraOperations operations;
+		private final @NonNull Converter<Object, Object> resultProcessing;
+
+		/* (non-Javadoc)
+		 * @see org.springframework.data.cassandra.repository.query.CassandraQueryExecution#execute(java.lang.String, java.lang.Class)
+		 */
+		@Override
+		public Object execute(String query, Class<?> type) {
+
+			return StreamUtils.createStreamFromIterator(operations.stream(query, type)).map(new Function<Object, Object>() {
+
+				@Override
+				public Object apply(Object t) {
+					return resultProcessing.convert(t);
+				}
+			});
+		}
+	}
+
 	/**
 	 * {@link CassandraQueryExecution} for collection returning queries.
 	 *
@@ -45,6 +75,9 @@ interface CassandraQueryExecution {
 
 		private final @NonNull CassandraOperations operations;
 
+		/* (non-Javadoc)
+		 * @see org.springframework.data.cassandra.repository.query.CassandraQueryExecution#execute(java.lang.String, java.lang.Class)
+		 */
 		@Override
 		public Object execute(String query, Class<?> type) {
 			return operations.select(query, type);
@@ -61,6 +94,9 @@ interface CassandraQueryExecution {
 
 		private final @NonNull CassandraOperations operations;
 
+		/* (non-Javadoc)
+		 * @see org.springframework.data.cassandra.repository.query.CassandraQueryExecution#execute(java.lang.String, java.lang.Class)
+		 */
 		@Override
 		public Object execute(String query, Class<?> type) {
 			return operations.selectOne(query, type);
@@ -77,6 +113,9 @@ interface CassandraQueryExecution {
 
 		private final @NonNull CassandraOperations operations;
 
+		/* (non-Javadoc)
+		 * @see org.springframework.data.cassandra.repository.query.CassandraQueryExecution#execute(java.lang.String, java.lang.Class)
+		 */
 		@Override
 		public Object execute(String query, Class<?> type) {
 			return operations.query(query);
@@ -94,6 +133,9 @@ interface CassandraQueryExecution {
 		private final @NonNull CassandraQueryExecution delegate;
 		private final @NonNull Converter<Object, Object> converter;
 
+		/* (non-Javadoc)
+		 * @see org.springframework.data.cassandra.repository.query.CassandraQueryExecution#execute(java.lang.String, java.lang.Class)
+		 */
 		@Override
 		public Object execute(String query, Class<?> type) {
 			return converter.convert(delegate.execute(query, type));
@@ -110,6 +152,9 @@ interface CassandraQueryExecution {
 
 		private final @NonNull ResultProcessor processor;
 
+		/* (non-Javadoc)
+		 * @see org.springframework.core.convert.converter.Converter#convert(java.lang.Object)
+		 */
 		@Override
 		public Object convert(Object source) {
 

spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/core/CassandraOperationsIntegrationTests.java

@@ -15,10 +15,12 @@
  */
 package org.springframework.data.cassandra.test.integration.core;
 
+import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.*;
 
 import java.util.ArrayList;
 import java.util.Date;
+import java.util.Iterator;
 import java.util.List;
 import java.util.UUID;
 
@@ -651,8 +653,8 @@ public class CassandraOperationsIntegrationTests extends AbstractSpringDataEmbed
 		log.debug("SingleSelect Book Title -> " + b.getTitle());
 		log.debug("SingleSelect Book Author -> " + b.getAuthor());
 
-		Assert.assertEquals(b.getTitle(), "Spring Data Cassandra Guide");
-		Assert.assertEquals(b.getAuthor(), "Cassandra Guru");
+		assertEquals(b.getTitle(), "Spring Data Cassandra Guide");
+		assertEquals(b.getAuthor(), "Cassandra Guru");
 
 	}
 
@@ -669,11 +671,11 @@ public class CassandraOperationsIntegrationTests extends AbstractSpringDataEmbed
 
 		log.debug("Book Count -> " + bookz.size());
 
-		Assert.assertEquals(bookz.size(), 20);
+		assertEquals(bookz.size(), 20);
 
 		for (Book b : bookz) {
 			Assert.assertTrue(b.isInStock());
-			Assert.assertEquals(BookCondition.NEW, b.getCondition());
+			assertEquals(BookCondition.NEW, b.getCondition());
 		}
 	}
 
@@ -685,7 +687,7 @@ public class CassandraOperationsIntegrationTests extends AbstractSpringDataEmbed
 
 		template.insert(books);
 
-		Assert.assertEquals(count, template.count(Book.class));
+		assertEquals(count, template.count(Book.class));
 	}
 
 
@@ -697,7 +699,25 @@ public class CassandraOperationsIntegrationTests extends AbstractSpringDataEmbed
 
 		template.insert(books);
 
-		Assert.assertEquals(count, template.count(Book.class));
+		assertEquals(count, template.count(Book.class));
 	}
 
+	/**
+	 * @see DATACASS-297
+	 */
+	@Test
+	public void stream() {
+
+		List<Book> books = getBookList(20);
+		template.insert(books);
+
+		Iterator<Book> iterator = template.stream("select * from book", Book.class);
+		List<Book> result = new ArrayList<Book>();
+		while (iterator.hasNext()) {
+			result.add(iterator.next());
+		}
+
+		assertThat(books.size(), is(20));
+		assertThat(books.get(0), is(instanceOf(Book.class)));
+	}
 }

spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/QueryIntegrationTests.java

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.cassandra.test.integration.repository.querymethods.declared;
 
+import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.*;
 
 import java.util.Collection;
@@ -22,6 +23,8 @@ import java.util.Date;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Consumer;
+import java.util.stream.Stream;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -35,6 +38,7 @@ import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 
 /**
  * @author Matthew T. Adams
+ * @author Mark Paluch
  */
 @RunWith(SpringJUnit4ClassRunner.class)
 public abstract class QueryIntegrationTests extends AbstractSpringDataEmbeddedCassandraIntegrationTest {
@@ -231,21 +235,6 @@ public abstract class QueryIntegrationTests extends AbstractSpringDataEmbeddedCa
 		assertEquals(saved.getNumberOfChildren(), value);
 	}
 
-	// TODO: @Test
-	// public void testUuidMethodResult() {
-	//
-	// Person saved = new Person();
-	// saved.setFirstname(uuid());
-	// saved.setLastname(uuid());
-	// saved.setUuid(UUID.randomUUID());
-	//
-	// saved = r.save(saved);
-	//
-	// UUID value = r.findSingleUuid(saved.getLastname(), saved.getFirstname());
-	//
-	// assertEquals(saved.getUuid(), value);
-	// }
-
 	@Test
 	public void testArrayMethodSingleResult() {
 
@@ -309,4 +298,31 @@ public abstract class QueryIntegrationTests extends AbstractSpringDataEmbeddedCa
 
 		assertFalse(optional.isPresent());
 	}
+
+	/**
+	 * @see DATACASS-297
+	 */
+	@Test
+	public void streamShouldReturnEntities() {
+
+		for (int i = 0; i < 100; i++) {
+		
+			Person person = new Person();
+			person.setFirstname(uuid());
+			person.setLastname(uuid());
+			person.setNumberOfChildren(i);
+			
+			personRepository.save(person);
+		}
+
+		Stream<Person> allPeople = personRepository.findAllPeople();
+		long count = allPeople.peek(new Consumer<Person>() {
+			@Override
+			public void accept(Person person) {
+				assertThat(person, is(instanceOf(Person.class)));
+			}
+		}).count();
+
+		assertThat(count, is(100L));
+	}
 }

spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/anno/PersonRepositoryWithQueryAnnotations.java

@@ -19,6 +19,7 @@ import java.util.Date;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Stream;
 
 import org.springframework.data.cassandra.repository.Query;
 import org.springframework.data.cassandra.test.integration.repository.querymethods.declared.Person;
@@ -73,4 +74,8 @@ public interface PersonRepositoryWithQueryAnnotations extends PersonRepository {
 	@Override
 	@Query("select * from person where lastname = ?0 and firstname = ?1")
 	Optional<Person> findOptionalWithLastnameAndFirstname(String last, String first);
+
+	@Override
+	@Query("select * from person")
+	Stream<Person> findAllPeople();
 }

spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/base/PersonRepository.java

@@ -19,6 +19,7 @@ import java.util.Date;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Stream;
 
 import org.springframework.data.cassandra.repository.CassandraRepository;
 import org.springframework.data.cassandra.test.integration.repository.querymethods.declared.Person;
@@ -28,6 +29,7 @@ import com.datastax.driver.core.ResultSet;
 
 /**
  * @author Matthew T. Adams
+ * @author Mark Paluch
  */
 @NoRepositoryBean
 public interface PersonRepository extends CassandraRepository<Person> {
@@ -52,4 +54,6 @@ public interface PersonRepository extends CassandraRepository<Person> {
 
 	Optional<Person> findOptionalWithLastnameAndFirstname(String last, String first);
 
+	Stream<Person> findAllPeople();
+
 }

spring-data-cassandra/src/test/resources/META-INF/PersonRepositoryWithNamedQueries.properties

@@ -8,4 +8,5 @@ Person.findSingleBirthdate=select birthdate from person where lastname = ?0 and
 Person.findSingleCool=select cool from person where lastname = ?0 and firstname = ?1
 Person.findSingleNumberOfChildren=select numberofchildren from person where lastname = ?0 and firstname = ?1
 Person.findOptionalWithLastnameAndFirstname=select * from person where lastname = ?0 and firstname = ?1
+Person.findAllPeople=select * from person