From 1b4eec9dd40141cd6359f6d6e8c94b93691eee3c Mon Sep 17 00:00:00 2001 From: Mark Paluch Date: Tue, 14 Jun 2016 16:26:56 +0200 Subject: [PATCH] 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. --- .../cassandra/core/CqlTemplate.java | 47 +++++++++++--- .../cassandra/core/CassandraOperations.java | 15 +++++ .../cassandra/core/CassandraTemplate.java | 62 ++++++++++++++++++- .../query/AbstractCassandraQuery.java | 7 ++- .../query/CassandraQueryExecution.java | 45 ++++++++++++++ .../CassandraOperationsIntegrationTests.java | 32 ++++++++-- .../declared/QueryIntegrationTests.java | 46 +++++++++----- .../PersonRepositoryWithQueryAnnotations.java | 5 ++ .../declared/base/PersonRepository.java | 4 ++ ...ersonRepositoryWithNamedQueries.properties | 1 + 10 files changed, 231 insertions(+), 33 deletions(-) diff --git a/spring-cql/src/main/java/org/springframework/cassandra/core/CqlTemplate.java b/spring-cql/src/main/java/org/springframework/cassandra/core/CqlTemplate.java index e9bf58586..c3b275fdc 100644 --- a/spring-cql/src/main/java/org/springframework/cassandra/core/CqlTemplate.java +++ b/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 diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraOperations.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraOperations.java index b306921ca..ea43ca596 100644 --- a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraOperations.java +++ b/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}. + *

+ * Returns a {@link java.util.Iterator} that wraps the a Cassandra {@link com.datastax.driver.core.ResultSet}. + * + * @param element return type + * @param query must not be empty and not {@literal null}. + * @param type must not be {@literal null}. + * @return + * @since 1.5 + */ + Iterator stream(String query, Class type); + /** * Execute query and convert ResultSet to the list of entities. * diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraTemplate.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraTemplate.java index a90c2cbf5..7440a4d88 100644 --- a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraTemplate.java +++ b/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 Iterator stream(final String query, Class type) { + + Assert.hasText(query, "Query must not be empty"); + Assert.notNull(type, "Type must not be null"); + + ResultSet resultSet = doExecute(new SessionCallback() { + + @Override + public ResultSet doInSession(Session s) throws DataAccessException { + return s.execute(query); + } + }); + + if (resultSet == null) { + return Collections.emptyList().iterator(); + } + + return new ResultSetIteratorAdapter(resultSet.iterator(), getExceptionTranslator(), new CassandraConverterRowCallback(cassandraConverter, type)); + } + protected List select(final Select query, CassandraConverterRowCallback readRowCallback) { ResultSet resultSet = doExecute(new SessionCallback() { @@ -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 implements Iterator{ + + private final Iterator iterator; + private final PersistenceExceptionTranslator exceptionTranslator; + private final CassandraConverterRowCallback rowCallback; + + public ResultSetIteratorAdapter(Iterator iterator, PersistenceExceptionTranslator exceptionTranslator, CassandraConverterRowCallback 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); + } + } + } } diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/AbstractCassandraQuery.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/AbstractCassandraQuery.java index d9e447152..ccb89c6e6 100644 --- a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/AbstractCassandraQuery.java +++ b/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 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 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 { diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/CassandraQueryExecution.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/CassandraQueryExecution.java index 1880f9e50..fe6ec3c43 100644 --- a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/repository/query/CassandraQueryExecution.java +++ b/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 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() { + + @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 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) { diff --git a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/core/CassandraOperationsIntegrationTests.java b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/core/CassandraOperationsIntegrationTests.java index 271c41fa4..7443ca5d9 100644 --- a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/core/CassandraOperationsIntegrationTests.java +++ b/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 books = getBookList(20); + template.insert(books); + + Iterator iterator = template.stream("select * from book", Book.class); + List result = new ArrayList(); + while (iterator.hasNext()) { + result.add(iterator.next()); + } + + assertThat(books.size(), is(20)); + assertThat(books.get(0), is(instanceOf(Book.class))); + } } diff --git a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/QueryIntegrationTests.java b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/QueryIntegrationTests.java index cb31400fc..ec169e251 100644 --- a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/QueryIntegrationTests.java +++ b/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 allPeople = personRepository.findAllPeople(); + long count = allPeople.peek(new Consumer() { + @Override + public void accept(Person person) { + assertThat(person, is(instanceOf(Person.class))); + } + }).count(); + + assertThat(count, is(100L)); + } } diff --git a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/anno/PersonRepositoryWithQueryAnnotations.java b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/anno/PersonRepositoryWithQueryAnnotations.java index a09b8d70f..a9a4b0dcb 100644 --- a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/anno/PersonRepositoryWithQueryAnnotations.java +++ b/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 findOptionalWithLastnameAndFirstname(String last, String first); + + @Override + @Query("select * from person") + Stream findAllPeople(); } diff --git a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/base/PersonRepository.java b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/base/PersonRepository.java index 26cb0eb36..a2d40526e 100644 --- a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/test/integration/repository/querymethods/declared/base/PersonRepository.java +++ b/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 { @@ -52,4 +54,6 @@ public interface PersonRepository extends CassandraRepository { Optional findOptionalWithLastnameAndFirstname(String last, String first); + Stream findAllPeople(); + } diff --git a/spring-data-cassandra/src/test/resources/META-INF/PersonRepositoryWithNamedQueries.properties b/spring-data-cassandra/src/test/resources/META-INF/PersonRepositoryWithNamedQueries.properties index 3c074f895..54109186f 100644 --- a/spring-data-cassandra/src/test/resources/META-INF/PersonRepositoryWithNamedQueries.properties +++ b/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