JSONResultConverter.java

package com.tinkerpop.rexster.gremlin.converter;

import com.tinkerpop.blueprints.Element;
import com.tinkerpop.blueprints.util.io.graphson.GraphSONMode;
import com.tinkerpop.blueprints.util.io.graphson.GraphSONUtility;
import com.tinkerpop.pipes.util.structures.Pair;
import com.tinkerpop.pipes.util.structures.Row;
import com.tinkerpop.pipes.util.structures.Table;
import com.tinkerpop.rexster.Tokens;
import org.codehaus.jettison.json.JSONArray;
import org.codehaus.jettison.json.JSONObject;

import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Converts a result from Gremlin to a JSONArray using GraphSON format.
 *
 * @author Blake Eggleston (bdeggleston.github.com)
 */
public class JSONResultConverter implements ResultConverter<JSONArray> {

    private final GraphSONMode mode;
    private final long offsetStart;
    private final long offsetEnd;
    private final Set<String> returnKeys;

    public JSONResultConverter(final GraphSONMode mode, final long offsetStart, final long offsetEnd,
                               final Set<String> returnKeys) {
        this.mode = mode;
        this.offsetEnd = offsetEnd;
        this.offsetStart = offsetStart;
        this.returnKeys = returnKeys;
    }

    public JSONArray convert(final Object result) throws Exception {
        return convert(result, false).getA();
    }

    public Pair<JSONArray, Long> convert(final Object result, final boolean returnTotal) throws Exception {
        JSONArray results = new JSONArray();
        long counter = 0;
        if (result == null) {
            // for example a script like g.clear()
            results = null;

        } else if (result instanceof Table) {
            final Table table = (Table) result;
            final Iterator<Row> rows = table.iterator();

            while (rows.hasNext()) {
                final Row row = rows.next();
                if (counter >= this.offsetStart && counter < this.offsetEnd)
                    results.put(prepareOutput(row));

                if (!returnTotal && counter >= this.offsetEnd)
                    break;

                counter++;
            }
        } else if (result instanceof Iterable) {
            for (Object o : (Iterable) result) {
                if (counter >= this.offsetStart && counter < this.offsetEnd)
                    results.put(prepareOutput(o));

                if (!returnTotal && counter >= this.offsetEnd)
                    break;

                counter++;
            }
        } else if (result instanceof Iterator) {
            final Iterator itty = (Iterator) result;
            while (itty.hasNext()) {
                Object current = itty.next();
                if (counter >= this.offsetStart && counter < this.offsetEnd)
                    results.put(prepareOutput(current));

                if (!returnTotal && counter >= this.offsetEnd)
                    break;

                counter++;
            }
        } else {
            results.put(prepareOutput(result));
        }

        return new Pair<JSONArray, Long>(results, counter);
    }

    private Object prepareOutput(final Object object) throws Exception {
        if (object == null) {
            return null;
        }
        if (object instanceof Element) {
            return GraphSONUtility.jsonFromElement((Element) object, returnKeys, this.mode);
        } else if (object instanceof Row) {
            final Row row = (Row) object;
            final List<String> columnNames = row.getColumnNames();
            final Map<String, Object> map = new HashMap<String, Object>();
            for (String columnName : columnNames) {
                map.put(columnName, prepareOutput(row.getColumn(columnName)));
            }

            return new JSONObject(map);
        } else if (object instanceof Map) {
            final JSONObject jsonObject = new JSONObject();
            final Map map = (Map) object;
            for (Object key : map.keySet()) {
                // force an error here by passing in a null key to the JSONObject.  That way a good error message
                // gets back to the user.
                if (key instanceof Element) {
                    final Element element = (Element) key;
                    final HashMap<String, Object> m = new HashMap<String, Object>();
                    m.put(Tokens._KEY, this.prepareOutput(element));
                    m.put(Tokens._VALUE, this.prepareOutput(map.get(key)));

                    jsonObject.put(element.getId().toString(), new JSONObject(m));
                } else {
                    jsonObject.put(key == null ? null : key.toString(), this.prepareOutput(map.get(key)));
                }
            }

            return jsonObject;
        } else if (object instanceof Table || object instanceof Iterable || object instanceof Iterator) {
            return this.convert(object);
        } else if (object instanceof Number || object instanceof Boolean) {
            return object;
        } else if (object == JSONObject.NULL) {
            return JSONObject.NULL;
        } else {
            return object.toString();
        }
    }
}