diff --git a/src/java/org/apache/cassandra/cql3/ResultSet.java b/src/java/org/apache/cassandra/cql3/ResultSet.java index 16f0d1b590..e9d2f6f8eb 100644 --- a/src/java/org/apache/cassandra/cql3/ResultSet.java +++ b/src/java/org/apache/cassandra/cql3/ResultSet.java @@ -32,6 +32,7 @@ import org.apache.cassandra.thrift.CqlResultType; import org.apache.cassandra.thrift.CqlRow; import org.apache.cassandra.utils.ByteBufferUtil; import org.apache.cassandra.service.pager.PagingState; +import com.google.common.annotations.VisibleForTesting; public class ResultSet { @@ -286,6 +287,12 @@ public class ResultSet names.add(name); } + @VisibleForTesting + public PagingState getPagingState() + { + return pagingState; + } + public void setHasMorePages(PagingState pagingState) { this.pagingState = pagingState; diff --git a/test/distributed/org/apache/cassandra/distributed/impl/Coordinator.java b/test/distributed/org/apache/cassandra/distributed/impl/Coordinator.java index 6553fb94de..345e7132e4 100644 --- a/test/distributed/org/apache/cassandra/distributed/impl/Coordinator.java +++ b/test/distributed/org/apache/cassandra/distributed/impl/Coordinator.java @@ -21,17 +21,23 @@ package org.apache.cassandra.distributed.impl; import java.net.InetSocketAddress; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.UUID; import java.util.concurrent.Future; +import com.google.common.collect.Iterators; + +import com.datastax.driver.core.ProtocolVersion; +import org.apache.cassandra.config.CFMetaData; import org.apache.cassandra.cql3.CQLStatement; import org.apache.cassandra.cql3.QueryOptions; import org.apache.cassandra.cql3.QueryProcessor; import org.apache.cassandra.cql3.UntypedResultSet; import org.apache.cassandra.cql3.statements.SelectStatement; +import org.apache.cassandra.db.Keyspace; import org.apache.cassandra.distributed.api.ConsistencyLevel; import org.apache.cassandra.distributed.api.ICoordinator; import org.apache.cassandra.distributed.api.IInstance; @@ -122,49 +128,64 @@ public class Coordinator implements ICoordinator throw new IllegalArgumentException("Page size should be strictly positive but was " + pageSize); return instance.sync(() -> { + ClientState clientState = makeFakeClientState(); ConsistencyLevel consistencyLevel = ConsistencyLevel.valueOf(consistencyLevelOrigin.name()); - CQLStatement prepared = QueryProcessor.getStatement(query, ClientState.forInternalCalls()).statement; + CQLStatement prepared = QueryProcessor.getStatement(query, clientState).statement; List boundBBValues = new ArrayList<>(); for (Object boundValue : boundValues) { boundBBValues.add(ByteBufferUtil.objectToBytes(boundValue)); } - prepared.validate(QueryState.forInternalCalls().getClientState()); + prepared.validate(clientState); assert prepared instanceof SelectStatement : "Only SELECT statements can be executed with paging"; - ClientState clientState = QueryState.forInternalCalls().getClientState(); SelectStatement selectStatement = (SelectStatement) prepared; - QueryOptions queryOptions = QueryOptions.create(toCassandraCL(consistencyLevel), - boundBBValues, - false, - pageSize, - null, - null, - Server.CURRENT_VERSION); - Pageable pageable = selectStatement.getPageableCommand(queryOptions); - // Usually pager fetches a single page (see SelectStatement#execute). We need to iterate over all - // of the results lazily. - QueryPager pager = QueryPagers.pager(pageable, toCassandraCL(consistencyLevel), clientState, null); - Iterator iter = RowUtil.toObjects(selectStatement.getResultMetadata().names, - UntypedResultSet.create(selectStatement, - pager, - pageSize).iterator()); + QueryState queryState = new QueryState(clientState); + QueryOptions initialOptions = QueryOptions.create(toCassandraCL(consistencyLevel), + boundBBValues, + false, + pageSize, + null, + null, + Server.CURRENT_VERSION); + + + ResultMessage.Rows initialRows = selectStatement.execute(queryState, initialOptions); + Iterator iter = new Iterator() { + ResultMessage.Rows rows = selectStatement.execute(queryState, initialOptions); + Iterator iter = RowUtil.toIter(rows); - // We have to make sure iterator is not running on main thread. - Iterator it = new Iterator() { public boolean hasNext() { - return instance.sync(() -> iter.hasNext()).call(); + if (iter.hasNext()) + return true; + + if (rows.result.metadata.getPagingState() == null) + return false; + + QueryOptions nextOptions = QueryOptions.create(toCassandraCL(consistencyLevel), + boundBBValues, + true, + pageSize, + rows.result.metadata.getPagingState(), + null, + Server.CURRENT_VERSION); + + rows = selectStatement.execute(queryState, nextOptions); + iter = Iterators.forArray(RowUtil.toObjects(initialRows.result.metadata.names, rows.result.rows)); + + return hasNext(); } public Object[] next() { - return instance.sync(() -> iter.next()).call(); + return iter.next(); } }; - return QueryResults.fromObjectArrayIterator(RowUtil.getColumnNames(selectStatement.getResultMetadata().names), it); + + return QueryResults.fromObjectArrayIterator(RowUtil.getColumnNames(initialRows.result.metadata.names), iter); }).call(); } diff --git a/test/distributed/org/apache/cassandra/distributed/impl/RowUtil.java b/test/distributed/org/apache/cassandra/distributed/impl/RowUtil.java index 50d501ec49..4fcb4f2e8a 100644 --- a/test/distributed/org/apache/cassandra/distributed/impl/RowUtil.java +++ b/test/distributed/org/apache/cassandra/distributed/impl/RowUtil.java @@ -19,10 +19,13 @@ package org.apache.cassandra.distributed.impl; import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; import java.util.Iterator; import java.util.List; import com.google.common.collect.Iterators; +import com.google.common.collect.Lists; import com.datastax.driver.core.ResultSet; import com.datastax.driver.core.Row; @@ -56,44 +59,7 @@ public class RowUtil public static Object[][] toObjects(ResultMessage.Rows rows) { - Object[][] result = new Object[rows.result.rows.size()][]; - List specs = rows.result.metadata.names; - for (int i = 0; i < rows.result.rows.size(); i++) - { - List row = rows.result.rows.get(i); - result[i] = new Object[row.size()]; - for (int j = 0; j < row.size(); j++) - { - ByteBuffer bb = row.get(j); - - if (bb != null) - result[i][j] = specs.get(j).type.getSerializer().deserialize(bb); - } - } - return result; - } - - public static Iterator toObjects(UntypedResultSet rs) - { - return toObjects(rs.metadata(), rs.iterator()); - } - - public static Iterator toObjects(List columnSpecs, Iterator rs) - { - return Iterators.transform(rs, - (row) -> { - Object[] objectRow = new Object[columnSpecs.size()]; - for (int i = 0; i < columnSpecs.size(); i++) - { - ColumnSpecification columnSpec = columnSpecs.get(i); - ByteBuffer bb = row.getBytes(columnSpec.name.toString()); - - if (bb != null) - objectRow[i] = columnSpec.type.getSerializer().deserialize(bb); - - } - return objectRow; - }); + return toObjects(rows.result.metadata.names, rows.result.rows); } public static Iterator toObjects(ResultSet rs) @@ -109,5 +75,64 @@ public class RowUtil }); } + public static Object[][] toObjects(List specs, List> rows) + { + Object[][] result = new Object[rows.size()][]; + for (int i = 0; i < rows.size(); i++) + { + List row = rows.get(i); + result[i] = new Object[row.size()]; + for (int j = 0; j < row.size(); j++) + { + ByteBuffer bb = row.get(j); + if (bb != null) + result[i][j] = specs.get(j).type.getSerializer().deserialize(bb); + } + } + return result; + } + + public static Iterator toIter(UntypedResultSet rs) + { + return toIter(rs.metadata(), rs.iterator()); + } + + public static Iterator toIter(ResultMessage.Rows rows) + { + return toIterInternal(rows.result.metadata.names, rows.result.rows); + } + + public static Iterator toIter(List columnSpecs, Iterator rs) + { + Iterator> iter = Iterators.transform(rs, + (row) -> { + List bbs = new ArrayList<>(columnSpecs.size()); + for (int i = 0; i < columnSpecs.size(); i++) + { + ColumnSpecification columnSpec = columnSpecs.get(i); + bbs.add(row.getBytes(columnSpec.name.toString())); + } + return bbs; + }); + return toIterInternal(columnSpecs, Lists.newArrayList(iter)); + } + + private static Iterator toIterInternal(List columnSpecs, List> rs) + { + return Iterators.transform(rs.iterator(), + (row) -> { + Object[] objectRow = new Object[columnSpecs.size()]; + for (int i = 0; i < columnSpecs.size(); i++) + { + ColumnSpecification columnSpec = columnSpecs.get(i); + ByteBuffer bb = row.get(i); + + if (bb != null) + objectRow[i] = columnSpec.type.getSerializer().deserialize(bb); + + } + return objectRow; + }); + } }