HttpBasedClient.java

/*
 * 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 com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.JsonParseException;

import org.apache.commons.lang3.StringUtils;
import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;

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

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;

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;

/**
 * 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.optJSONObject(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 : String.join(",", indices);
    final String typs = types == null ? null : String.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 = getTotal(result);
        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"));
  }

  /**
   * Returns total hits returned by Elasticsearch.
   *
   * The implementation constists of two ways to retrieve total.
   *
   * The first one, happy path, it trying to get it from an object
   * "total" which has a field called "value".
   *
   * The second one is trying to get that number directly from hits
   * object under "total" value as it is not an object.
   *
   * The first way seems to be present in all new versions
   * in ES (for sure e.g. 7.x). By this way we also
   * support backward compatibility with older ES version
   * when it was not providing "total" as object yet.
   *
   * @param result result from search
   * @return number of total hits
   */
  private long getTotal(HttpResponse<JsonNode> result) {
    final JSONObject hitsObject = result.getBody().getObject().getJSONObject("hits");
    try {
      return hitsObject.getJSONObject("total").getLong("value");
    } catch (final JSONException ex) {
      return hitsObject.getLong("total");
    }
  }

  @Override
  public void close() {
  }

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