column.java

package basho.riak;

import com.basho.riak.client.RiakClient;
import com.basho.riak.client.operations.FetchDatatype;
import com.basho.riak.client.operations.FetchSet;
import com.basho.riak.client.operations.UpdateSet;
import com.basho.riak.client.operations.datatypes.SetUpdate;
import com.basho.riak.client.operations.indexes.BinIndexQuery;
import com.basho.riak.client.operations.kv.DeleteValue;
import com.basho.riak.client.operations.kv.FetchValue;
import com.basho.riak.client.operations.kv.StoreValue;
import com.basho.riak.client.query.Location;
import com.basho.riak.client.query.Namespace;
import com.basho.riak.client.query.RiakObject;
import com.basho.riak.client.query.indexes.StringBinIndex;
import com.basho.riak.client.util.BinaryValue;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.util.*;

public class Column<T> {
    public String columnName;
    private Namespace bucket;
    public RiakClient client;
    private Location setLocation;

    public Column(String columnName, RiakClient client) {
        this.columnName = columnName;
        this.client = client;
        this.bucket = new Namespace(columnName);
        this.setLocation = new Location(new Namespace("sets", columnName), "keys");
    }

    // Basic object fetch on the basis of its key
    private Map getColumnObject(String key) throws Exception {
        Location loc = new Location(bucket, key);
        FetchValue fetch = new FetchValue.Builder(loc).build();
        return client.execute(fetch).getValue(Map.class);
    }

    // Outputs the entire column as JSON, i.e. {"column":"<name>","objects":[...]}
    public String asJson() throws Exception {
        HashMap<String, Object> map = new HashMap<>();
        List<Map> objects = getObjects();
        map.put("objects", objects);
        map.put("column", columnName);
        return new ObjectMapper().writeValueAsString(map);
    }

    // Retrieves the column's keys on the basis of 2i
    public Set<String> getKeys() throws Exception {
        Set<String> builder = new HashSet<>();
        Location loc = setLocation;
        FetchSet fetch = new FetchSet.Builder(loc).build();
        client.execute(fetch).getDatatype().viewAsSet().forEach((BinaryValue bin) -> {
           builder.add(bin.toString());
        });
        return builder;
    }

    // Assembles a list of objects on the basis of the getKeys() operation
    public List<Map> getObjects() throws Exception {
        Set<String> keys = getKeys();
        List<Map> builder = new ArrayList<Map>();
        for (String key : keys) {
            Map obj = getColumnObject(key);
            builder.add(obj);
        }
        return builder;
    }

    // Assembles a list of values for a single field for all objects in the column
    public List<Object> getField(String field)
            throws Exception {
        List<Map> maps = getObjects();
        List<Object> objects = new ArrayList<Object>();

        maps.forEach((Map m) -> {
            if (m.containsKey(field)) {
                objects.add(m.get(field));
            }
        });

        return objects;
    }

    // Adds an object to the column. Must be a Map.
    public void addObject(String key, Map object) throws Exception {
        Location loc = new Location(bucket, key);
        String json = new ObjectMapper().writeValueAsString(object);
        RiakObject obj = new RiakObject()
                .setContentType("application/json")
                .setValue(BinaryValue.create(json));
        StoreValue store = new StoreValue.Builder(obj)
                .withLocation(loc)
                .build();
        SetUpdate addKey = new SetUpdate()
                .add(BinaryValue.create(key));
        UpdateSet update = new UpdateSet.Builder(setLocation, addKey).build();
        client.execute(store);
        client.execute(update);
    }

    // Removes an object from the column on the basis of its key
    public void removeObject(String key) throws Exception {
        Location loc = new Location(bucket, key);
        DeleteValue delete = new DeleteValue.Builder(loc).build();
        client.execute(delete);
    }

    // Returns a list of objects that meet a field-level criterion, e.g. all objects
    // for which the field "foo" equals ["bar", "baq", "baz"]
    public List<Map> selectObjects(String field, Object comparison) throws Exception {
        List<Map> builder = new ArrayList<>();
        List<Map> maps = getObjects();

        maps.forEach((Map m) -> {
            if (m.containsKey(field) && m.get(field).equals(comparison)) {
                builder.add(m);
            }
        });

        return builder;
    }
}