Examples with SimpleTokenBackedResultsPage com.palantir.util.paging.SimpleTokenBackedResultsPage used on opensource projects

Example 1 with SimpleTokenBackedResultsPage

use of com.palantir.util.paging.SimpleTokenBackedResultsPage in project atlasdb by palantir.

the class CqlKeyValueService method getRangeWithPageCreator.

public <T> ClosableIterator<RowResult<T>> getRangeWithPageCreator(TableReference tableRef, RangeRequest rangeRequest, long timestamp, com.datastax.driver.core.ConsistencyLevel consistency, Supplier<ResultsExtractor<T>> resultsExtractor) {
    if (rangeRequest.isReverse()) {
        throw new UnsupportedOperationException();
    }
    if (rangeRequest.isEmptyRange()) {
        return ClosableIterators.wrap(ImmutableList.<RowResult<T>>of().iterator());
    }
    final int batchHint = rangeRequest.getBatchHint() == null ? 100 : rangeRequest.getBatchHint();
    final ColumnSelection selection = rangeRequest.getColumnNames().isEmpty() ? ColumnSelection.all() : ColumnSelection.create(rangeRequest.getColumnNames());
    final byte[] endExclusive = rangeRequest.getEndExclusive();
    final StringBuilder bindQuery = new StringBuilder();
    bindQuery.append("SELECT * FROM " + getFullTableName(tableRef) + " WHERE token(" + fieldNameProvider.row() + ") >= token(?) ");
    if (endExclusive.length > 0) {
        bindQuery.append("AND token(" + fieldNameProvider.row() + ") < token(?) ");
    }
    bindQuery.append("LIMIT " + batchHint);
    final String getLastRowQuery = "SELECT * FROM " + getFullTableName(tableRef) + " WHERE " + fieldNameProvider.row() + " = ?";
    return ClosableIterators.wrap(new AbstractPagingIterable<RowResult<T>, TokenBackedBasicResultsPage<RowResult<T>, byte[]>>() {

        @Override
        protected TokenBackedBasicResultsPage<RowResult<T>, byte[]> getFirstPage() throws Exception {
            return getPage(rangeRequest.getStartInclusive());
        }

        @Override
        protected TokenBackedBasicResultsPage<RowResult<T>, byte[]> getNextPage(TokenBackedBasicResultsPage<RowResult<T>, byte[]> previous) throws Exception {
            return getPage(previous.getTokenForNextPage());
        }

        TokenBackedBasicResultsPage<RowResult<T>, byte[]> getPage(final byte[] startKey) throws Exception {
            BoundStatement boundStatement = getPreparedStatement(tableRef, bindQuery.toString(), session).setConsistencyLevel(consistency).bind();
            boundStatement.setBytes(0, ByteBuffer.wrap(startKey));
            if (endExclusive.length > 0) {
                boundStatement.setBytes(1, ByteBuffer.wrap(endExclusive));
            }
            ResultSet resultSet = session.execute(boundStatement);
            List<Row> rows = Lists.newArrayList(resultSet.all());
            cqlKeyValueServices.logTracedQuery(bindQuery.toString(), resultSet, session, cqlStatementCache.normalQuery);
            byte[] maxRow = null;
            ResultsExtractor<T> extractor = resultsExtractor.get();
            for (Row row : rows) {
                byte[] rowName = getRowName(row);
                if (maxRow == null) {
                    maxRow = rowName;
                } else {
                    maxRow = PtBytes.BYTES_COMPARATOR.max(maxRow, rowName);
                }
            }
            if (maxRow == null) {
                return new SimpleTokenBackedResultsPage<>(endExclusive, ImmutableList.of(), false);
            }
            // get the rest of the last row
            BoundStatement boundLastRow = getPreparedStatement(tableRef, getLastRowQuery, session).bind();
            boundLastRow.setBytes(fieldNameProvider.row(), ByteBuffer.wrap(maxRow));
            try {
                resultSet = session.execute(boundLastRow);
            } catch (com.datastax.driver.core.exceptions.UnavailableException e) {
                throw new InsufficientConsistencyException("This operation requires all Cassandra" + " nodes to be up and available.", e);
            }
            rows.addAll(resultSet.all());
            cqlKeyValueServices.logTracedQuery(getLastRowQuery, resultSet, session, cqlStatementCache.normalQuery);
            for (Row row : rows) {
                extractor.internalExtractResult(timestamp, selection, getRowName(row), getColName(row), getValue(row), getTs(row));
            }
            SortedMap<byte[], SortedMap<byte[], T>> resultsByRow = Cells.breakCellsUpByRow(extractor.asMap());
            return ResultsExtractor.getRowResults(endExclusive, maxRow, resultsByRow);
        }
    }.iterator());
}