com.netflix.astyanax.thrift.ThriftAllRowsQueryImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of astyanax-thrift Show documentation
The newest version!
package com.netflix.astyanax.thrift;

import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicReference;

import org.apache.cassandra.thrift.ColumnParent;
import org.apache.cassandra.thrift.KeyRange;
import org.apache.cassandra.thrift.KeySlice;
import org.apache.cassandra.thrift.SlicePredicate;
import org.apache.cassandra.thrift.SliceRange;
import org.apache.cassandra.thrift.Cassandra.Client;
import org.apache.cassandra.utils.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Function;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.util.concurrent.ListenableFuture;
import com.netflix.astyanax.CassandraOperationType;
import com.netflix.astyanax.ExceptionCallback;
import com.netflix.astyanax.RowCallback;
import com.netflix.astyanax.connectionpool.ConnectionContext;
import com.netflix.astyanax.connectionpool.Host;
import com.netflix.astyanax.connectionpool.OperationResult;
import com.netflix.astyanax.connectionpool.TokenRange;
import com.netflix.astyanax.connectionpool.exceptions.ConnectionException;
import com.netflix.astyanax.connectionpool.impl.OperationResultImpl;
import com.netflix.astyanax.model.ByteBufferRange;
import com.netflix.astyanax.model.ColumnFamily;
import com.netflix.astyanax.model.ColumnSlice;
import com.netflix.astyanax.model.Rows;
import com.netflix.astyanax.partitioner.Partitioner;
import com.netflix.astyanax.query.AllRowsQuery;
import com.netflix.astyanax.query.CheckpointManager;
import com.netflix.astyanax.shallows.EmptyCheckpointManager;
import com.netflix.astyanax.thrift.model.ThriftRowsSliceImpl;

public class ThriftAllRowsQueryImpl implements AllRowsQuery {
    private final static Logger LOG = LoggerFactory.getLogger(ThriftAllRowsQueryImpl.class);
    
    private final ThriftColumnFamilyQueryImpl query;
    protected SlicePredicate predicate = new SlicePredicate().setSlice_range(ThriftUtils.createAllInclusiveSliceRange());
    protected CheckpointManager checkpointManager = new EmptyCheckpointManager();
    
    protected ColumnFamily columnFamily;
    private ExceptionCallback exceptionCallback;
    private int     blockSize       = 100;
    private boolean repeatLastToken = true;
    private Integer nThreads;
    private String  startToken      ;
    private String  endToken        ;
    private Boolean includeEmptyRows;
    
    public ThriftAllRowsQueryImpl(ThriftColumnFamilyQueryImpl query) {
        this.columnFamily = query.columnFamily;
        this.query = query;
    }
    
    protected List getNextBlock(final KeyRange range) {
        ThriftKeyspaceImpl keyspace = query.keyspace;
        
        while (true) {
            try {
                return keyspace.connectionPool.executeWithFailover(
                        new AbstractKeyspaceOperationImpl>(
                                keyspace.tracerFactory.newTracer(CassandraOperationType.GET_ROWS_RANGE, columnFamily),
                                query.pinnedHost, keyspace.getKeyspaceName()) {
                            @Override
                            public List internalExecute(Client client, ConnectionContext context)
                                    throws Exception {
                                
                                List slice = client.get_range_slices(
                                        new ColumnParent().setColumn_family(columnFamily.getName()), predicate,
                                        range, ThriftConverter.ToThriftConsistencyLevel(query.consistencyLevel));
                                
                                return slice;
                            }

                            @Override
                            public ByteBuffer getRowKey() {
                                if (range.getStart_key() != null)
                                    return range.start_key;
                                return null;
                            }
                        }, query.retry).getResult();
            }
            catch (ConnectionException e) {
                // Let exception callback handle this exception. If it
                // returns false then
                // we return an empty result which the iterator's
                // hasNext() to return false.
                // If no exception handler is provided then simply
                // return an empty set as if the
                // there is no more data
                if (this.getExceptionCallback() == null) {
                    throw new RuntimeException(e);
                }
                else {
                    if (!this.getExceptionCallback().onException(e)) {
                        return new ArrayList();
                    }
                }
            }
        }
    }

    @Override
    public OperationResult> execute() throws ConnectionException {
        return new OperationResultImpl>(Host.NO_HOST, 
                new ThriftAllRowsImpl(query.keyspace.getPartitioner(), this, columnFamily), 0);
    }

    @Override
    public ListenableFuture>> executeAsync() throws ConnectionException {
        throw new UnsupportedOperationException("executeAsync not supported here.  Use execute()");
    }

    private boolean shouldIgnoreEmptyRows() {
        if (getIncludeEmptyRows() == null) {
            if (getPredicate().isSetSlice_range() && getPredicate().getSlice_range().getCount() == 0) {
                return false;
            }
        }
        else {
            return !getIncludeEmptyRows();
        }

        return true;
    }

    @Override
    public void executeWithCallback(final RowCallback callback) throws ConnectionException {
        final ThriftKeyspaceImpl keyspace = query.keyspace;
        final Partitioner partitioner = keyspace.getPartitioner();
        final AtomicReference error = new AtomicReference();
        final boolean bIgnoreTombstones = shouldIgnoreEmptyRows();

        List> ranges;
        if (this.getConcurrencyLevel() != null) {
            ranges = Lists.newArrayList();
            int nThreads = this.getConcurrencyLevel();
            List tokens = partitioner.splitTokenRange(
                    startToken == null ? partitioner.getMinToken() : startToken, 
                    endToken == null   ? partitioner.getMaxToken() : endToken, 
                    nThreads);
            for (TokenRange range : tokens) {
                try {
                    String currentToken = checkpointManager.getCheckpoint(range.getStartToken());
                    if (currentToken == null) {
                        currentToken = range.getStartToken();
                    }
                    else if (currentToken.equals(range.getEndToken())) {
                        continue;
                    }
                    ranges.add(Pair.create(currentToken, range.getEndToken()));
                } catch (Exception e) {
                    throw ThriftConverter.ToConnectionPoolException(e);
                }
            }
        }
        else {
            ranges = Lists.transform(keyspace.describeRing(true), new Function> () {
                @Override
                public Pair apply(TokenRange input) {
                    return Pair.create(input.getStartToken(), input.getEndToken());
                }
            });
        }
        final CountDownLatch doneSignal = new CountDownLatch(ranges.size());

        for (final Pair tokenPair : ranges) {
            // Prepare the range of tokens for this token range
            final KeyRange range = new KeyRange()
                    .setCount(getBlockSize())
                    .setStart_token(tokenPair.left)
                    .setEnd_token(tokenPair.right);

            query.executor.submit(new Callable() {
                private boolean firstBlock = true;
                
                @Override
                public Void call() throws Exception {
                    if (error.get() == null && internalRun()) {
                        query.executor.submit(this);
                    }
                    else {
                        doneSignal.countDown();
                    }
                    return null;
                }

                private boolean internalRun() throws Exception {
                    try {
                        // Get the next block
                        List ks = keyspace.connectionPool.executeWithFailover(
                                new AbstractKeyspaceOperationImpl>(keyspace.tracerFactory
                                        .newTracer(CassandraOperationType.GET_ROWS_RANGE,
                                                columnFamily), query.pinnedHost, keyspace
                                        .getKeyspaceName()) {
                                    @Override
                                    public List internalExecute(Client client, ConnectionContext context)
                                            throws Exception {
                                        return client.get_range_slices(new ColumnParent()
                                                .setColumn_family(columnFamily.getName()),
                                                predicate, range, ThriftConverter
                                                        .ToThriftConsistencyLevel(query.consistencyLevel));
                                    }

                                    @Override
                                    public ByteBuffer getRowKey() {
                                        if (range.getStart_key() != null)
                                            return ByteBuffer.wrap(range.getStart_key());
                                        return null;
                                    }
                                }, query.retry.duplicate()).getResult();

                        // Notify the callback
                        if (!ks.isEmpty()) {
                            KeySlice lastRow = Iterables.getLast(ks);
                            boolean bContinue = (ks.size() == getBlockSize());

                            if (getRepeatLastToken()) {
                                if (firstBlock) {
                                    firstBlock = false;
                                }
                                else {
                                    ks.remove(0);
                                }
                            }
                            
                            if (bIgnoreTombstones) {
                                Iterator iter = ks.iterator();
                                while (iter.hasNext()) {
                                    if (iter.next().getColumnsSize() == 0)
                                        iter.remove();
                                }
                            }
                            Rows rows = new ThriftRowsSliceImpl(ks, columnFamily
                                    .getKeySerializer(), columnFamily.getColumnSerializer());
                            try {
                                callback.success(rows);
                            }
                            catch (Throwable t) {
                                ConnectionException ce = ThriftConverter.ToConnectionPoolException(t);
                                error.set(ce);
                                return false;
                            }
                            
                            if (bContinue) {
                                // Determine the start token for the next page
                                String token = partitioner.getTokenForKey(lastRow.bufferForKey()).toString();
                                checkpointManager.trackCheckpoint(tokenPair.left, token);
                                if (getRepeatLastToken()) {
                                    range.setStart_token(partitioner.getTokenMinusOne(token));
                                }
                                else {
                                    range.setStart_token(token);
                                }
                            }
                            else {
                                checkpointManager.trackCheckpoint(tokenPair.left, tokenPair.right);
                                return false;
                            }
                        }
                        else {
                            checkpointManager.trackCheckpoint(tokenPair.left, tokenPair.right);
                            return false;
                        }
                    }
                    catch (Exception e) {
                        ConnectionException ce = ThriftConverter.ToConnectionPoolException(e);
                        if (!callback.failure(ce)) {
                            error.set(ce);
                            return false;
                        }
                    }

                    return true;
                }
            });
        }
        // Block until all threads finish
        try {
            doneSignal.await();
        }
        catch (InterruptedException e) {
            LOG.debug("Execution interrupted on get all rows for keyspace " + keyspace.getKeyspaceName());
        }

        if (error.get() != null) {
            throw error.get();
        }
    }
    
    public AllRowsQuery setExceptionCallback(ExceptionCallback cb) {
        exceptionCallback = cb;
        return this;
    }

    protected ExceptionCallback getExceptionCallback() {
        return this.exceptionCallback;
    }

    @Override
    public AllRowsQuery setThreadCount(int numberOfThreads) {
        setConcurrencyLevel(numberOfThreads);
        return this;
    }
    
    @Override
    public AllRowsQuery setConcurrencyLevel(int numberOfThreads) {
        this.nThreads = numberOfThreads;
        return this;
    }


    @Override
    public AllRowsQuery setCheckpointManager(CheckpointManager manager) {
        this.checkpointManager = manager;
        return this;
    }

    @Override
    public AllRowsQuery withColumnSlice(C... columns) {
        if (columns != null)
            predicate.setColumn_names(columnFamily.getColumnSerializer().toBytesList(Arrays.asList(columns)))
                    .setSlice_rangeIsSet(false);
        return this;
    }

    @Override
    public AllRowsQuery withColumnSlice(Collection columns) {
        if (columns != null)
            predicate.setColumn_names(columnFamily.getColumnSerializer().toBytesList(columns)).setSlice_rangeIsSet(
                    false);
        return this;
    }

    @Override
    public AllRowsQuery withColumnRange(C startColumn, C endColumn, boolean reversed, int count) {
        predicate.setSlice_range(ThriftUtils.createSliceRange(columnFamily.getColumnSerializer(), startColumn,
                endColumn, reversed, count));
        return this;
    }

    @Override
    public AllRowsQuery withColumnRange(ByteBuffer startColumn, ByteBuffer endColumn, boolean reversed, int count) {
        predicate.setSlice_range(new SliceRange(startColumn, endColumn, reversed, count));
        return this;
    }

    @Override
    public AllRowsQuery withColumnSlice(ColumnSlice slice) {
        if (slice.getColumns() != null) {
            predicate.setColumn_names(columnFamily.getColumnSerializer().toBytesList(slice.getColumns()))
                    .setSlice_rangeIsSet(false);
        }
        else {
            predicate.setSlice_range(ThriftUtils.createSliceRange(columnFamily.getColumnSerializer(),
                    slice.getStartColumn(), slice.getEndColumn(), slice.getReversed(), slice.getLimit()));
        }
        return this;
    }

    @Override
    public AllRowsQuery withColumnRange(ByteBufferRange range) {
        predicate.setSlice_range(new SliceRange().setStart(range.getStart()).setFinish(range.getEnd())
                .setCount(range.getLimit()).setReversed(range.isReversed()));
        return this;
    }

    @Override
    public AllRowsQuery setBlockSize(int blockSize) {
        return setRowLimit(blockSize);
    }

    @Override
    public AllRowsQuery setRowLimit(int rowLimit) {
        this.blockSize = rowLimit;
        return this;
    }

    public int getBlockSize() {
        return blockSize;
    }

    @Override
    public AllRowsQuery setRepeatLastToken(boolean repeatLastToken) {
        this.repeatLastToken = repeatLastToken;
        return this;
    }

    public boolean getRepeatLastToken() {
        return this.repeatLastToken;
    }

    protected Integer getConcurrencyLevel() {
        return this.nThreads;
    }
    
    public AllRowsQuery setIncludeEmptyRows(boolean flag) {
        this.includeEmptyRows = flag;
        return this;
    }

    public String getStartToken() {
        return this.startToken;
    }
    
    public String getEndToken() {
        return this.endToken;
    }
    
    @Override
    public AllRowsQuery forTokenRange(BigInteger startToken, BigInteger endToken) {
        return forTokenRange(startToken.toString(), endToken.toString());
    }
    
    public AllRowsQuery forTokenRange(String startToken, String endToken) {
        this.startToken = startToken;
        this.endToken = endToken;
        return this;
    }
    
    SlicePredicate getPredicate() {
        return predicate;
    }
    
    Boolean getIncludeEmptyRows() {
        return this.includeEmptyRows;
    }
}