/*
* 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 + "]";
}
}