org.apache.zeppelin.elasticsearch.client.HttpBasedClient.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.zeppelin.elasticsearch.client.HttpBasedClient.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.zeppelin.elasticsearch.client;

import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.util.Iterator;
import java.util.Map;
import java.util.Properties;

import org.apache.commons.lang3.StringUtils;
import org.apache.zeppelin.elasticsearch.ElasticsearchInterpreter;
import org.apache.zeppelin.elasticsearch.action.ActionException;
import org.apache.zeppelin.elasticsearch.action.ActionResponse;
import org.apache.zeppelin.elasticsearch.action.AggWrapper;
import org.apache.zeppelin.elasticsearch.action.AggWrapper.AggregationType;
import org.apache.zeppelin.elasticsearch.action.HitWrapper;
import org.json.JSONArray;
import org.json.JSONObject;

import com.google.common.base.Joiner;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.JsonParseException;
import com.mashape.unirest.http.HttpResponse;
import com.mashape.unirest.http.JsonNode;
import com.mashape.unirest.http.Unirest;
import com.mashape.unirest.http.exceptions.UnirestException;
import com.mashape.unirest.request.HttpRequest;
import com.mashape.unirest.request.HttpRequestWithBody;

/**
 * Elasticsearch client using the HTTP API.
 */
public class HttpBasedClient implements ElasticsearchClient {

    private static final String QUERY_STRING_TEMPLATE = "{ \"query\": { \"query_string\": { \"query\": \"_Q_\", \"analyze_wildcard\": \"true\" } } }";

    private final String host;
    private final int port;
    private final String username;
    private final String password;

    private final Gson gson = new GsonBuilder().setPrettyPrinting().create();

    public HttpBasedClient(Properties props) {
        this.host = props.getProperty(ElasticsearchInterpreter.ELASTICSEARCH_HOST);
        this.port = Integer.parseInt(props.getProperty(ElasticsearchInterpreter.ELASTICSEARCH_PORT));
        this.username = props.getProperty(ElasticsearchInterpreter.ELASTICSEARCH_BASIC_AUTH_USERNAME);
        this.password = props.getProperty(ElasticsearchInterpreter.ELASTICSEARCH_BASIC_AUTH_PASSWORD);
    }

    private boolean isSucceeded(HttpResponse response) {
        return response.getStatus() >= 200 && response.getStatus() < 300;
    }

    private JSONObject getParentField(JSONObject parent, String[] fields) {
        JSONObject obj = parent;
        for (int i = 0; i < fields.length - 1; i++) {
            obj = obj.getJSONObject(fields[i]);
        }
        return obj;
    }

    private JSONArray getFieldAsArray(JSONObject obj, String field) {
        final String[] fields = field.split("/");
        final JSONObject parent = getParentField(obj, fields);
        return parent.getJSONArray(fields[fields.length - 1]);
    }

    private String getFieldAsString(HttpResponse<JsonNode> response, String field) {
        return getFieldAsString(response.getBody(), field);
    }

    private String getFieldAsString(JsonNode json, String field) {
        return json.getObject().get(field).toString();
    }

    private long getFieldAsLong(HttpResponse<JsonNode> response, String field) {
        final String[] fields = field.split("/");
        final JSONObject obj = getParentField(response.getBody().getObject(), fields);
        return obj.getLong(fields[fields.length - 1]);
    }

    private String getUrl(String index, String type, String id, boolean useSearch) {
        try {
            final StringBuilder buffer = new StringBuilder();
            buffer.append("http://").append(host).append(":").append(port).append("/");
            if (StringUtils.isNotEmpty(index)) {
                buffer.append(index);

                if (StringUtils.isNotEmpty(type)) {
                    buffer.append("/").append(type);

                    if (StringUtils.isNotEmpty(id)) {
                        if (useSearch) {
                            final String encodedId = URLEncoder.encode(id, "UTF-8");
                            if (id.equals(encodedId)) {
                                // No difference, use directly the id
                                buffer.append("/").append(id);
                            } else {
                                // There are differences: to avoid problems with some special characters
                                // such as / and # in id, use a "terms" query
                                buffer.append("/_search?source=").append(URLEncoder
                                        .encode("{\"query\":{\"terms\":{\"_id\":[\"" + id + "\"]}}}", "UTF-8"));
                            }
                        } else {
                            buffer.append("/").append(id);
                        }
                    }
                }
            }
            return buffer.toString();
        } catch (final UnsupportedEncodingException e) {
            throw new ActionException(e);
        }
    }

    private String getUrl(String[] indices, String[] types) {
        final String inds = indices == null ? null : Joiner.on(",").join(indices);
        final String typs = types == null ? null : Joiner.on(",").join(types);
        return getUrl(inds, typs, null, false);
    }

    @Override
    public ActionResponse get(String index, String type, String id) {
        ActionResponse response = null;
        try {
            final HttpRequest request = Unirest.get(getUrl(index, type, id, true));
            if (StringUtils.isNotEmpty(username)) {
                request.basicAuth(username, password);
            }

            final HttpResponse<String> result = request.asString();
            final boolean isSucceeded = isSucceeded(result);

            if (isSucceeded) {
                final JsonNode body = new JsonNode(result.getBody());
                if (body.getObject().has("_index")) {
                    response = new ActionResponse().succeeded(true)
                            .hit(new HitWrapper(getFieldAsString(body, "_index"), getFieldAsString(body, "_type"),
                                    getFieldAsString(body, "_id"), getFieldAsString(body, "_source")));
                } else {
                    final JSONArray hits = getFieldAsArray(body.getObject(), "hits/hits");
                    final JSONObject hit = (JSONObject) hits.iterator().next();
                    response = new ActionResponse().succeeded(true).hit(new HitWrapper(hit.getString("_index"),
                            hit.getString("_type"), hit.getString("_id"), hit.opt("_source").toString()));
                }
            } else {
                if (result.getStatus() == 404) {
                    response = new ActionResponse().succeeded(false);
                } else {
                    throw new ActionException(result.getBody());
                }
            }
        } catch (final UnirestException e) {
            throw new ActionException(e);
        }
        return response;
    }

    @Override
    public ActionResponse delete(String index, String type, String id) {
        ActionResponse response = null;
        try {
            final HttpRequest request = Unirest.delete(getUrl(index, type, id, true));
            if (StringUtils.isNotEmpty(username)) {
                request.basicAuth(username, password);
            }

            final HttpResponse<String> result = request.asString();
            final boolean isSucceeded = isSucceeded(result);

            if (isSucceeded) {
                final JsonNode body = new JsonNode(result.getBody());
                response = new ActionResponse().succeeded(true).hit(new HitWrapper(getFieldAsString(body, "_index"),
                        getFieldAsString(body, "_type"), getFieldAsString(body, "_id"), null));
            } else {
                throw new ActionException(result.getBody());
            }
        } catch (final UnirestException e) {
            throw new ActionException(e);
        }
        return response;
    }

    @Override
    public ActionResponse index(String index, String type, String id, String data) {
        ActionResponse response = null;
        try {
            HttpRequestWithBody request = null;
            if (StringUtils.isEmpty(id)) {
                request = Unirest.post(getUrl(index, type, id, false));
            } else {
                request = Unirest.put(getUrl(index, type, id, false));
            }
            request.header("Accept", "application/json").header("Content-Type", "application/json").body(data)
                    .getHttpRequest();
            if (StringUtils.isNotEmpty(username)) {
                request.basicAuth(username, password);
            }

            final HttpResponse<JsonNode> result = request.asJson();
            final boolean isSucceeded = isSucceeded(result);

            if (isSucceeded) {
                response = new ActionResponse().succeeded(true)
                        .hit(new HitWrapper(getFieldAsString(result, "_index"), getFieldAsString(result, "_type"),
                                getFieldAsString(result, "_id"), null));
            } else {
                throw new ActionException(result.getBody().toString());
            }
        } catch (final UnirestException e) {
            throw new ActionException(e);
        }
        return response;
    }

    @Override
    public ActionResponse search(String[] indices, String[] types, String query, int size) {
        ActionResponse response = null;

        if (!StringUtils.isEmpty(query)) {
            // The query can be either JSON-formatted, nor a Lucene query
            // So, try to parse as a JSON => if there is an error, consider the query a Lucene one
            try {
                gson.fromJson(query, Map.class);
            } catch (final JsonParseException e) {
                // This is not a JSON (or maybe not well formatted...)
                query = QUERY_STRING_TEMPLATE.replace("_Q_", query);
            }
        }

        try {
            final HttpRequestWithBody request = Unirest.post(getUrl(indices, types) + "/_search?size=" + size)
                    .header("Content-Type", "application/json");

            if (StringUtils.isNoneEmpty(query)) {
                request.header("Accept", "application/json").body(query);
            }

            if (StringUtils.isNotEmpty(username)) {
                request.basicAuth(username, password);
            }

            final HttpResponse<JsonNode> result = request.asJson();
            final JSONObject body = result.getBody() != null ? result.getBody().getObject() : null;

            if (isSucceeded(result)) {
                final long total = getFieldAsLong(result, "hits/total");

                response = new ActionResponse().succeeded(true).totalHits(total);

                if (containsAggs(result)) {
                    JSONObject aggregationsMap = body.getJSONObject("aggregations");
                    if (aggregationsMap == null) {
                        aggregationsMap = body.getJSONObject("aggs");
                    }

                    for (final String key : aggregationsMap.keySet()) {
                        final JSONObject aggResult = aggregationsMap.getJSONObject(key);
                        if (aggResult.has("buckets")) {
                            // Multi-bucket aggregations
                            final Iterator<Object> buckets = aggResult.getJSONArray("buckets").iterator();
                            while (buckets.hasNext()) {
                                response.addAggregation(
                                        new AggWrapper(AggregationType.MULTI_BUCKETS, buckets.next().toString()));
                            }
                        } else {
                            response.addAggregation(
                                    new AggWrapper(AggregationType.SIMPLE, aggregationsMap.toString()));
                        }
                        break; // Keep only one aggregation
                    }
                } else if (size > 0 && total > 0) {
                    final JSONArray hits = getFieldAsArray(body, "hits/hits");
                    final Iterator<Object> iter = hits.iterator();

                    while (iter.hasNext()) {
                        final JSONObject hit = (JSONObject) iter.next();
                        final Object data = hit.opt("_source") != null ? hit.opt("_source") : hit.opt("fields");
                        response.addHit(new HitWrapper(hit.getString("_index"), hit.getString("_type"),
                                hit.getString("_id"), data.toString()));
                    }
                }
            } else {
                throw new ActionException(body.get("error").toString());
            }
        } catch (final UnirestException e) {
            throw new ActionException(e);
        }

        return response;
    }

    private boolean containsAggs(HttpResponse<JsonNode> result) {
        return result.getBody() != null
                && (result.getBody().getObject().has("aggregations") || result.getBody().getObject().has("aggs"));
    }

    @Override
    public void close() {
    }

    @Override
    public String toString() {
        return "HttpBasedClient [host=" + host + ", port=" + port + ", username=" + username + "]";
    }
}