/* * 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.IOException; import java.net.InetAddress; import java.net.UnknownHostException; import java.util.HashMap; import java.util.HashSet; import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Properties; import java.util.Set; import org.apache.commons.lang3.StringUtils; import org.apache.zeppelin.elasticsearch.ElasticsearchInterpreter; import org.apache.zeppelin.elasticsearch.action.ActionResponse; import org.apache.zeppelin.elasticsearch.action.AggWrapper; import org.apache.zeppelin.elasticsearch.action.HitWrapper; import org.elasticsearch.action.delete.DeleteResponse; import org.elasticsearch.action.get.GetResponse; import org.elasticsearch.action.index.IndexResponse; import org.elasticsearch.action.search.SearchAction; import org.elasticsearch.action.search.SearchRequestBuilder; import org.elasticsearch.action.search.SearchResponse; import org.elasticsearch.client.Client; import org.elasticsearch.client.transport.TransportClient; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.transport.InetSocketTransportAddress; import org.elasticsearch.common.xcontent.XContentBuilder; import org.elasticsearch.common.xcontent.XContentFactory; import org.elasticsearch.common.xcontent.XContentHelper; import org.elasticsearch.index.query.QueryBuilders; import org.elasticsearch.search.SearchHit; import org.elasticsearch.search.SearchHitField; import org.elasticsearch.search.aggregations.Aggregation; import org.elasticsearch.search.aggregations.Aggregations; import org.elasticsearch.search.aggregations.InternalMultiBucketAggregation; import org.elasticsearch.search.aggregations.bucket.InternalSingleBucketAggregation; import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation; import org.elasticsearch.search.aggregations.metrics.InternalMetricsAggregation; import com.google.gson.Gson; import com.google.gson.GsonBuilder; import com.google.gson.JsonSyntaxException; /** * Elasticsearch client using the transport protocol. */ public class TransportBasedClient implements ElasticsearchClient { private final Gson gson = new GsonBuilder().setPrettyPrinting().create(); private final Client client; public TransportBasedClient(Properties props) throws UnknownHostException { final String host = props.getProperty(ElasticsearchInterpreter.ELASTICSEARCH_HOST); final int port = Integer.parseInt( props.getProperty(ElasticsearchInterpreter.ELASTICSEARCH_PORT)); final String clusterName = props.getProperty(ElasticsearchInterpreter.ELASTICSEARCH_CLUSTER_NAME); final Settings settings = Settings.settingsBuilder() .put("cluster.name", clusterName) .put(props) .build(); client = TransportClient.builder().settings(settings).build() .addTransportAddress(new InetSocketTransportAddress(InetAddress.getByName(host), port)); } @Override public ActionResponse get(String index, String type, String id) { final GetResponse getResp = client .prepareGet(index, type, id) .get(); return new ActionResponse() .succeeded(getResp.isExists()) .hit(new HitWrapper( getResp.getIndex(), getResp.getType(), getResp.getId(), getResp.getSourceAsString())); } @Override public ActionResponse delete(String index, String type, String id) { final DeleteResponse delResp = client .prepareDelete(index, type, id) .get(); return new ActionResponse() .succeeded(delResp.isFound()) .hit(new HitWrapper( delResp.getIndex(), delResp.getType(), delResp.getId(), null)); } @Override public ActionResponse index(String index, String type, String id, String data) { final IndexResponse idxResp = client .prepareIndex(index, type, id) .setSource(data) .get(); return new ActionResponse() .succeeded(idxResp.isCreated()) .hit(new HitWrapper( idxResp.getIndex(), idxResp.getType(), idxResp.getId(), null)); } @Override public ActionResponse search(String[] indices, String[] types, String query, int size) { final SearchRequestBuilder reqBuilder = new SearchRequestBuilder( client, SearchAction.INSTANCE); reqBuilder.setIndices(); if (indices != null) { reqBuilder.setIndices(indices); } if (types != null) { reqBuilder.setTypes(types); } 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 { @SuppressWarnings("rawtypes") final Map source = gson.fromJson(query, Map.class); reqBuilder.setExtraSource(source); } catch (final JsonSyntaxException e) { // This is not a JSON (or maybe not well formatted...) reqBuilder.setQuery(QueryBuilders.queryStringQuery(query).analyzeWildcard(true)); } } reqBuilder.setSize(size); final SearchResponse searchResp = reqBuilder.get(); final ActionResponse actionResp = new ActionResponse() .succeeded(true) .totalHits(searchResp.getHits().getTotalHits()); if (searchResp.getAggregations() != null) { setAggregations(searchResp.getAggregations(), actionResp); } else { for (final SearchHit hit: searchResp.getHits()) { // Fields can be found either in _source, or in fields (it depends on the query) // => specific for elasticsearch's version < 5 // String src = hit.getSourceAsString(); if (src == null) { final Map<String, Object> hitFields = new HashMap<>(); for (final SearchHitField hitField : hit.getFields().values()) { hitFields.put(hitField.getName(), hitField.getValues()); } src = gson.toJson(hitFields); } actionResp.addHit(new HitWrapper(hit.getIndex(), hit.getType(), hit.getId(), src)); } } return actionResp; } private void setAggregations(Aggregations aggregations, ActionResponse actionResp) { // Only the result of the first aggregation is returned // final Aggregation agg = aggregations.asList().get(0); if (agg instanceof InternalMetricsAggregation) { actionResp.addAggregation(new AggWrapper(AggWrapper.AggregationType.SIMPLE, XContentHelper.toString((InternalMetricsAggregation) agg).toString())); } else if (agg instanceof InternalSingleBucketAggregation) { actionResp.addAggregation(new AggWrapper(AggWrapper.AggregationType.SIMPLE, XContentHelper.toString((InternalSingleBucketAggregation) agg).toString())); } else if (agg instanceof InternalMultiBucketAggregation) { final Set<String> headerKeys = new HashSet<>(); final List<Map<String, Object>> buckets = new LinkedList<>(); final InternalMultiBucketAggregation multiBucketAgg = (InternalMultiBucketAggregation) agg; for (final MultiBucketsAggregation.Bucket bucket : multiBucketAgg.getBuckets()) { try { final XContentBuilder builder = XContentFactory.jsonBuilder(); bucket.toXContent(builder, null); actionResp.addAggregation( new AggWrapper(AggWrapper.AggregationType.MULTI_BUCKETS, builder.string())); } catch (final IOException e) { // Ignored } } } } @Override public void close() { if (client != null) { client.close(); } } @Override public String toString() { return "TransportBasedClient []"; } }