Spring/spring-data-cassandra
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

DATACASS-809 - Add queryForStream methods to CqlOperations.

Browse Source 
CqlOperations now exposes queryForStream methods to consume a query as java.util.stream.Stream accepting a RowMapper.
This commit is contained in:
Mark Paluch
2020-10-01 10:45:18 +02:00
parent 4f332dc5ac
commit d6a32814ca
8 changed files with 254 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -19,7 +19,6 @@ import java.util.List;
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 class CassandraTemplate implements CassandraOperations, ApplicationEventP
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 @@ public class CassandraTemplate implements CassandraOperations, ApplicationEventP
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)

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

@@ -752,8 +752,7 @@ public class AsyncCqlTemplate extends CassandraAccessor implements AsyncCqlOpera
}
/**
* 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 @@ public class AsyncCqlTemplate extends CassandraAccessor implements AsyncCqlOpera
}
/**
* 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}.

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

@@ -18,16 +18,17 @@ package org.springframework.data.cassandra.core.cql;
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 @@ public interface CqlOperations {
*/
<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 @@ public interface CqlOperations {
<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 @@ public interface CqlOperations {
<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
// -------------------------------------------------------------------------

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

@@ -18,7 +18,11 @@ package org.springframework.data.cassandra.core.cql;
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 class CqlTemplate extends CassandraAccessor implements CqlOperations {
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 class CqlTemplate extends CassandraAccessor implements CqlOperations {
*/
@Override
public Map<String, Object> queryForMap(Statement<?> statement) throws DataAccessException {
// noinspection ConstantConditions
return queryForObject(statement, newColumnMapRowMapper());
}
@@ -485,6 +498,17 @@ public class CqlTemplate extends CassandraAccessor implements CqlOperations {
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 class CqlTemplate extends CassandraAccessor implements CqlOperations {
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 class CqlTemplate extends CassandraAccessor implements CqlOperations {
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 @@ public class CqlTemplate extends CassandraAccessor implements CqlOperations {
}
/**
* 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 @@ public class CqlTemplate extends CassandraAccessor implements CqlOperations {
}
/**
* 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 @@ public class CqlTemplate extends CassandraAccessor implements CqlOperations {
}
/**
* 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 @@ public class CqlTemplate extends CassandraAccessor implements CqlOperations {
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 @@ public class CqlTemplate extends CassandraAccessor implements CqlOperations {
@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);
}
}
}

8
spring-data-cassandra/src/main/kotlin/org/springframework/data/cassandra/core/cql/CqlOperationsExtensions.kt
View File

@@ -18,6 +18,7 @@ package org.springframework.data.cassandra.core.cql
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 java.util.stream.Stream
import kotlin.reflect.KClass
/**
@@ -120,3 +121,10 @@ fun CqlOperations.query(cql: String, vararg args: Any, function: (Row) -> Unit):
*/
fun <T : Any> CqlOperations.query(cql: String, vararg args: Any, function: (Row, Int) -> T): List<T> =
query(cql, RowMapper { row, i -> function(row, i) }, *args)
/**
* Extension for [queryForStream.query] providing a RowMapper-like function
* variant: `query("...", arg1, argN){ row, i -> }`.
*/
fun <T : Any> CqlOperations.queryForStream(cql: String, vararg args: Any, function: (Row, Int) -> T): Stream<T> =
queryForStream(cql, RowMapper { row, i -> function(row, i) }, *args)

35
spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/cql/CqlTemplateUnitTests.java
View File

@@ -24,6 +24,8 @@ import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -337,6 +339,19 @@ class CqlTemplateUnitTests {
});
}
@Test // DATACASS-809
public void queryForStreamWithResultSetStatementExtractorWithArgumentsShouldCallExecution() {
doTestStrings(5, ConsistencyLevel.ONE, null, "foo", cqlTemplate -> {
Stream<String> result = cqlTemplate.queryForStream(SimpleStatement.newInstance("SELECT * from USERS"),
(row, index) -> row.getString(0));
assertThat(result.collect(Collectors.toList())).hasSize(3).contains("Walter", "Hank", " Jesse");
verify(session).execute(any(Statement.class));
});
}
@Test // DATACASS-292
void queryStatementShouldTranslateExceptions() {
@@ -610,6 +625,21 @@ class CqlTemplateUnitTests {
verify(preparedStatement).bind("a", "b");
}
@Test // DATACASS-809
public void queryForStreanPreparedStatementCreatorAndBinderAndMapperShouldReturnResult() {
when(session.execute(boundStatement)).thenReturn(resultSet);
when(resultSet.spliterator()).thenReturn(Collections.singleton(row).spliterator());
Stream<Row> rows = template.queryForStream(session -> preparedStatement, ps -> {
ps.bind("a", "b");
return boundStatement;
}, (row, rowNum) -> row);
assertThat(rows).hasSize(1).contains(row);
verify(preparedStatement).bind("a", "b");
}
@Test // DATACASS-292
void queryForObjectPreparedStatementShouldBeEmpty() {
@@ -728,9 +758,10 @@ class CqlTemplateUnitTests {
String[] results = { "Walter", "Hank", " Jesse" };
List<Row> rows = Arrays.asList(row, row, row);
when(this.session.execute((Statement) any())).thenReturn(resultSet);
when(this.resultSet.iterator()).thenReturn(Arrays.asList(row, row, row).iterator());
when(this.resultSet.spliterator()).thenCallRealMethod();
when(this.resultSet.iterator()).thenReturn(rows.iterator());
when(this.resultSet.spliterator()).thenReturn(rows.spliterator());
when(this.row.getString(0)).thenReturn(results[0], results[1], results[2]);
when(this.session.prepare(anyString())).thenReturn(preparedStatement);

7
spring-data-cassandra/src/test/kotlin/org/springframework/data/cassandra/core/cql/CqlOperationsExtensionsUnitTests.kt
View File

@@ -136,4 +136,11 @@ class CqlOperationsExtensionsUnitTests {
operations.query("", 3) { row, _ -> row.columnDefinitions }
verify { operations.query(eq(""), any<RowMapper<Person>>(), eq(3)) }
}
@Test // DATACASS-809
fun `queryForStream(String, RowMapper, array) extension should call its Java counterpart`() {
operations.queryForStream("", 3) { row, _ -> row.columnDefinitions }
verify { operations.queryForStream(eq(""), any<RowMapper<Person>>(), eq(3)) }
}
}

1
src/main/asciidoc/new-features.adoc
View File

@@ -9,6 +9,7 @@ This chapter summarizes changes and new features for each release.
* <<cassandra.auditing,Reactive auditing>> enabled through `@EnableReactiveCassandraAuditing`. `@EnableCassandraAuditing` no longer registers `ReactiveAuditingEntityCallback`.
* Reactive SpEL support in `@Query` query methods.
* Configuration of the keyspace per `Statement` through `CqlTemplate` and `QueryOptions`.
* Revised `CqlOperations` with new `queryForStream(…)` methods returning a `Stream` with transparent pagination.
[[new-features.3-0-0]]
== What's new in Spring Data for Apache Cassandra 3.0