JestRichResult.java

package com.codetroopers.play.elasticsearch.jest;

import com.codetroopers.play.elasticsearch.Index;
import com.codetroopers.play.elasticsearch.IndexUtils;
import com.google.common.base.Objects;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.google.gson.*;
import io.searchbox.client.JestResult;
import io.searchbox.core.search.facet.*;

import javax.annotation.Nullable;
import javax.validation.constraints.NotNull;
import java.lang.reflect.Constructor;
import java.util.List;
import java.util.Map;

/**
 * @author cgatay
 */
public class JestRichResult {
    private static final Gson gson = createGsonWithDateFormat();
    private static final Map<String, Class<? extends Facet>> FACET_TYPE_CLASSES;
    private JestResult result;

    static {
        ImmutableMap.Builder<String, Class<? extends Facet>> tmp = ImmutableMap.builder();
        tmp.put(QueryFacet.TYPE.toLowerCase(), QueryFacet.class);
        tmp.put(StatisticalFacet.TYPE.toLowerCase(), StatisticalFacet.class);
        tmp.put(TermsStatsFacet.TYPE.toLowerCase(), TermsStatsFacet.class);
        tmp.put(DateHistogramFacet.TYPE.toLowerCase(), DateHistogramFacet.class);
        tmp.put(GeoDistanceFacet.TYPE.toLowerCase(), GeoDistanceFacet.class);
        tmp.put(FilterFacet.TYPE.toLowerCase(), FilterFacet.class);
        tmp.put(TermsFacet.TYPE.toLowerCase(), TermsFacet.class);
        tmp.put(RangeFacet.TYPE.toLowerCase(), RangeFacet.class);
        tmp.put(HistogramFacet.TYPE.toLowerCase(), HistogramFacet.class);
        FACET_TYPE_CLASSES = tmp.build();
    }

    public JestRichResult(@Nullable JestResult result) {
        this.result = result;
    }

    public List<Facet> getFacets() {
        List<Facet> out = Lists.newArrayList();
        final JsonObject jsonObject = safeResult().getJsonObject();
        if (jsonObject != null) {
            final JsonElement facets = jsonObject.get("facets");
            if (facets != null) {
                for (Map.Entry<String, JsonElement> facetEntry : ((JsonObject) facets).entrySet()) {
                    JsonObject facet = facetEntry.getValue().getAsJsonObject();
                    final JsonElement jsonElement = facet.get("_type");
                    if (jsonElement != null) {
                        final Class<? extends Facet> facetClazz = FACET_TYPE_CLASSES.get(jsonElement.getAsString().toLowerCase());

                        try {
                            final Constructor<? extends Facet> constructor = facetClazz.getConstructor(String.class, JsonObject.class);
                            out.add(constructor.newInstance(facetEntry.getKey(), facetEntry.getValue()));
                        } catch (Exception e) {
                            e.printStackTrace();
                        }

                    }
                }
            }
        }
        return out;
    }

    public int getTotalHits() {
        final JsonElement hits = safeResult().getJsonObject().get("hits");
        if (hits != null) {
            final JsonElement total = hits.getAsJsonObject().get("total");
            if (total != null) {
                return total.getAsInt();
            }
        }
        return 0;
    }

    public List<Result> getHits() {
        List<Result> out = Lists.newArrayList();
        final JsonElement hits = safeResult().getJsonObject().get("hits");
        if (hits != null) {
            final JsonElement hitsList = hits.getAsJsonObject().get("hits");
            if (hitsList != null) {
                final JsonArray jsonArray = hitsList.getAsJsonArray();
                for (JsonElement jsonElement : jsonArray) {
                    out.add(new Result(jsonElement.getAsJsonObject()));
                }
            }
        }
        return out;
    }

    @Nullable
    public <T extends Index> T getFirstHit(Class<T> clazz) {
        if (safeResult().isSucceeded()) {
            return new Result(safeResult().getJsonObject()).getObject(clazz);
        }
        return null;
    }

    @Nullable
    public String getId() {
        final JsonObject jsonObject = safeResult().getJsonObject();
        if (jsonObject.has("_id")) {
            return jsonObject.get("_id").getAsString();
        }
        return null;
    }

    @NotNull
    public JestResult getResult() {
        return safeResult();
    }


    public String getPathToResult() {
        return safeResult().getPathToResult();
    }

    public String getJsonString() {
        return safeResult().getJsonString();
    }

    public Object getValue(String key) {
        return safeResult().getValue(key);
    }

    public <T extends Facet> List<T> getFacets(Class<T> type) {
        return safeResult().getFacets(type);
    }

    public boolean isSucceeded() {
        return safeResult().isSucceeded();
    }

    public <T> List<T> getSourceAsObjectList(Class<T> type) {
        return safeResult().getSourceAsObjectList(type);
    }

    public String getErrorMessage() {
        return safeResult().getErrorMessage();
    }

    public JsonObject getJsonObject() {
        return safeResult().getJsonObject();
    }

    public <T> T getSourceAsObject(Class<T> clazz) {
        return safeResult().getSourceAsObject(clazz);
    }

    /**
     * Allows elasticSearch and Gson to talk the same language regarding dates (ISO8601 format)
     */
    static Gson createGsonWithDateFormat() {
        return new GsonBuilder()
                .setDateFormat("yyyy-MM-dd'T'HH:mm:ss'Z'").create();
    }

    JestResult safeResult() {
        if (result == null) {
            result = new JestResult(gson);
            safeResult().setJsonObject(new JsonObject());
            safeResult().setSucceeded(false);
        }
        return result;
    }

    public static class Result {
        private final String index;
        private final String type;
        private final String id;
        private Float score;
        private final JsonObject result;
        private final JsonObject source;
        private transient Index convertedObject = null;

        public Result(JsonObject resultLine) {
            this.result = resultLine;
            this.index = resultLine.get("_index").getAsString();
            this.type = resultLine.get("_type").getAsString();
            this.id = resultLine.get("_id").getAsString();
            final JsonElement scoreElement = resultLine.get("_score");
            try {
                this.score = resultLine.has("_score") ? scoreElement.getAsFloat() : 1.0f;
            } catch (Exception e) {
                this.score = 1.0f;
            }
            this.source = resultLine.has("_source") ? resultLine.get("_source").getAsJsonObject() : null;
        }

        public String index() {
            return index;
        }

        public String type() {
            return type;
        }

        public String id() {
            return id;
        }

        public Float score() {
            return score;
        }

        public JsonObject result() {
            return result;
        }

        @SuppressWarnings("unchecked")
        public Map<String, Object> sourceAsMap() {
            return gson.fromJson(source, Map.class);
        }

        public JsonObject source() {
            return source;
        }

        @SuppressWarnings("unchecked")
        public <T extends Index> T getObject(Class<T> clazz) {
            if (convertedObject == null) {
                // Get Data Map
                Map<String, Object> map = sourceAsMap();
                // Create a new Indexable Object for the return
                T objectIndexable = IndexUtils.getInstanceIndex(clazz);
                T t = (T) objectIndexable.fromIndex(map);
                t.id = id();
                convertedObject = t;
            }
            return (T) convertedObject;
        }

        @Override
        public String toString() {
            return Objects.toStringHelper(this)
                    .add("source", source)
                    .add("score", score)
                    .add("id", id)
                    .add("type", type)
                    .add("index", index)
                    .toString();
        }
    }
}