com.netflix.astyanax.recipes.functions.RowCopierFunction Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of astyanax-recipes Show documentation
astyanax-recipes
There is a newer version: 3.10.2
package com.netflix.astyanax.recipes.functions;

import java.io.Flushable;
import java.io.IOException;
import java.util.Set;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Function;
import com.google.common.collect.Sets;
import com.netflix.astyanax.ColumnListMutation;
import com.netflix.astyanax.Keyspace;
import com.netflix.astyanax.MutationBatch;
import com.netflix.astyanax.connectionpool.exceptions.ConnectionException;
import com.netflix.astyanax.model.Column;
import com.netflix.astyanax.model.ColumnFamily;
import com.netflix.astyanax.model.Row;

/**
 * Function to copy rows into a target column family
 * 
 * TODO:  Failover, retry
 *  
 * @author elandau
 *
 * @param 
 * @param 
 */
public class RowCopierFunction implements Function, Boolean>, Flushable {
    private static final Logger LOG = LoggerFactory.getLogger(RowCopierFunction.class);
    
    private static final int DEFAULT_BATCH_SIZE = 100;
    
    public static class Builder {
        private final ColumnFamily columnFamily;
        private final Keyspace          keyspace;
        private       int               batchSize           = DEFAULT_BATCH_SIZE;
        
        public Builder(Keyspace keyspace, ColumnFamily columnFamily) {
            this.columnFamily = columnFamily;
            this.keyspace     = keyspace;
        }
        
        public Builder withBatchSize(int batchSize) {
            this.batchSize = batchSize;
            return this;
        }
        
        public RowCopierFunction build() {
            return new RowCopierFunction(this);
        }
    }

    public static  Builder builder(Keyspace keyspace, ColumnFamily columnFamily) {
        return new Builder(keyspace, columnFamily);
    }
    
    private final ColumnFamily columnFamily;
    private final Keyspace          keyspace;
    private final int               batchSize;
    private final ThreadLocal context = new ThreadLocal();
    private final Set contexts = Sets.newIdentityHashSet();
    
    private static class ThreadContext {
        MutationBatch mb;
        int counter = 0;
    }
    
    private RowCopierFunction(Builder builder) {
        this.columnFamily = builder.columnFamily;
        this.batchSize    = builder.batchSize;
        this.keyspace     = builder.keyspace;
    }
    
    @Override
    public Boolean apply(Row row) {
        ThreadContext context = this.context.get();
        if (context == null) {
            context = new ThreadContext();
            context.mb = keyspace.prepareMutationBatch();
            this.context.set(context);
            
            synchronized (this) {
                contexts.add(context);
            }
        }
        
        ColumnListMutation mbRow = context.mb.withRow(columnFamily, row.getKey());
        context.mb.lockCurrentTimestamp();
        for (Column column : row.getColumns()) {
            mbRow.setTimestamp(column.getTimestamp());
            mbRow.putColumn(column.getName(), column.getByteBufferValue(), column.getTtl());
        }
        
        context.counter++;
        if (context.counter == batchSize) {
            try {
                context.mb.execute();
                context.counter = 0;
            }
            catch (Exception e) {
                LOG.error("Failed to write mutation", e);
                return false;
            }
            
        }
        return true;
    }

    @Override
    public void flush() throws IOException {
        for (ThreadContext context : contexts) {
            try {
                context.mb.execute();
            } catch (ConnectionException e) {
                LOG.error("Failed to write mutation", e);
            }
        }
    }

}