MutationUtils.java

package com.fullcontact.hadoop.cassandra;

import com.google.common.collect.Lists;
import org.apache.cassandra.db.marshal.AbstractType;
import org.apache.cassandra.db.marshal.CompositeType;
import org.apache.cassandra.db.marshal.UTF8Type;
import org.apache.cassandra.thrift.Column;
import org.apache.cassandra.thrift.ColumnOrSuperColumn;
import org.apache.cassandra.thrift.Mutation;
import org.apache.cassandra.utils.ByteBufferUtil;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.List;

/**
 * Useful utilities for CQL3 manipulation using Cassandra Thrift interfaces
 *
 * @author Michael Rose <michael@fullcontact.com>
 */
public class MutationUtils {
    /**
     * Gets a mutation
     *
     * e.x.
     * mutationList.add(getMutation(INDEXED_VALUE, mapper.writeValueAsBytes(identity)));
     * mutationList.add(getMutation(TIMESTAMP_KEY, ByteBufferUtil.bytes(System.currentTimeMillis())));
     *
     *
     * @param column col name
     * @param value col value (will be turned into a byte array)
     * @return Thrift Mutation
     * @throws IOException
     */
    public static Mutation getMutation(String column, String value) throws IOException {
        return getMutation(column, ByteBufferUtil.bytes(value));
    }

    public static Mutation getMutation(String column, byte[] value) throws IOException {
        return getMutation(column, ByteBuffer.wrap(value));
    }

    public static Mutation getMutation(String column, ByteBuffer value) throws IOException {
        ByteBuffer columnName = getCompositeColumnName("2", column);

        Column c = new Column();
        c.setName(columnName);
        c.setValue(value);
        c.setTimestamp(System.currentTimeMillis());

        Mutation m = new Mutation();
        m.setColumn_or_supercolumn(new ColumnOrSuperColumn());
        m.column_or_supercolumn.setColumn(c);

        return m;
    }

    /**
     * Useful for CQL composite keys, e.g. a rowkey of (account, email) could be found with
     * getCompositeColumnName("112", "test@test.com")
     * 
     */
    public static ByteBuffer getCompositeColumnName(String... parts) {
        return getCompositeColumnName(Lists.newArrayList(parts));
    }

    public static ByteBuffer getCompositeColumnName(List<String> parts) {
        List<AbstractType<?>> keyTypes = Lists.newArrayListWithExpectedSize(parts.size());

        for (String ignored : parts) {
            keyTypes.add(UTF8Type.instance);
        }

        CompositeType compositeKey = CompositeType.getInstance(keyTypes);
        CompositeType.Builder builder = new CompositeType.Builder(compositeKey);

        for (String part : parts) {
            builder.add(ByteBufferUtil.bytes(part));
        }

        return builder.build();
    }
}