elasticsearch high level客户端实现

import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.support.WriteRequest.RefreshPolicy;
import org.elasticsearch.action.support.master.AcknowledgedResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.client.indices.GetIndexRequest;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.index.reindex.UpdateByQueryRequest;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
import org.elasticsearch.search.sort.SortOrder;

import kd.bos.exception.BosErrorCode;
import kd.bos.exception.KDException;
import kd.bos.fulltext.common.ConstantKeys;

/**
 * Rest High Level 方式与  Elastic Search 交互
 * 
 * Default port: 9200
 * 
 * index/indice: 相当于mysql中的data base
 * type: 相当于mysql中的表
 * document: 相当于mysql中的行
 * field: 相当于mysql中的字段
 * 
 * 若与服务端版本不一致可能出现版本问题
 * @author rd_jianbin_lai
 */
public class HighLevelAccess {
	
	private HighLevelAccess(){
		
	}
	
	private static final String RETRY_ON_CONFLICT_KEY = "fulltext.retryonconflict";
	
	public static RestHighLevelClient createRestHighLevelClient(String ip, int port, String username, String password) {
		//use username and password for authentication
		if(username != null && password != null) {
			final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
			credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(username, password));
			return new RestHighLevelClient(RestClient.builder(new HttpHost(ip, port))
					.setHttpClientConfigCallback(new RestClientBuilder.HttpClientConfigCallback() {
						public HttpAsyncClientBuilder customizeHttpClient(HttpAsyncClientBuilder httpClientBuilder) {
							httpClientBuilder.disableAuthCaching();
							return httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
						}
					}));
		}
		//without authentication
		return new RestHighLevelClient(RestClient.builder(new HttpHost(ip, port, "http")));
	}

public static boolean isIndexExist(RestHighLevelClient client, String index) throws IOException {
		GetIndexRequest req = new GetIndexRequest(index);
		return client.indices().exists(req, RequestOptions.DEFAULT);
	}

	public static void createIndex(RestHighLevelClient client, String index, Map<String, Integer> settings)
			throws IOException {
		CreateIndexRequest req = new CreateIndexRequest(index);
		client.indices().create(req.settings(settings), RequestOptions.DEFAULT);
	}

	public static void createIndexMapping(RestHighLevelClient client, String index, String type,String fieldsMapping) throws IOException {
		PutMappingRequest req = new PutMappingRequest(index);
		req.type(type);
		req.source(fieldsMapping, XContentType.JSON);
		AcknowledgedResponse res = client.indices().putMapping(req, RequestOptions.DEFAULT);
		if (!res.isAcknowledged()) {
			throw new KDException(BosErrorCode.fulltextException, "Failed to create index:" + index  + res.toString());
		}
	}
	
	public static void updateIndexSettings(RestHighLevelClient client, String[] indexs,
			Map<String, Integer> indexSettings) throws IOException {
		String shardsNum = "number_of_shards";
		if (indexSettings.containsKey(shardsNum)) {
			indexSettings.remove(shardsNum);
		}
		UpdateSettingsRequest req = new UpdateSettingsRequest(indexs);
		req.settings(indexSettings);
		AcknowledgedResponse res = client.indices().putSettings(req, RequestOptions.DEFAULT);
		if (!res.isAcknowledged()) {
			throw new KDException(BosErrorCode.fulltextException,  "failed to update indexs settings:" + res.toString());
		}
	}

public static List<BulkItemResponse> batchSave(RestHighLevelClient client, String index, String type,
			List<Map<String, Object>> listMap) throws IOException {
		BulkRequest req = new BulkRequest();
		req.setRefreshPolicy(RefreshPolicy.IMMEDIATE);
		for (Map<String, Object> map : listMap) {
			Object id = map.get(ConstantKeys.ESID_FIELD);
			map.remove(ConstantKeys.ESID_FIELD, id);
			IndexRequest indexRequest = new IndexRequest(index, type, String.valueOf(id)).source(map);
			UpdateRequest updateRequest = new UpdateRequest(index, type, String.valueOf(id))
					.upsert(indexRequest);
			updateRequest.doc(indexRequest);
			updateRequest.retryOnConflict(Integer.parseInt(System.getProperty(RETRY_ON_CONFLICT_KEY, "5")));
			req.add(updateRequest);
		}
		BulkResponse res = client.bulk(req, RequestOptions.DEFAULT);
		List<BulkItemResponse> failureRes = new ArrayList<>();
		if (res.hasFailures()) {
			for (BulkItemResponse resp : res.getItems()) {
				if (resp.isFailed()) {
					failureRes.add(resp);
				}
			}
		}
		return failureRes;
	}

public static long updateByQuery(RestHighLevelClient client, String index, String type,
			Map<String, Object> filterData, Map<String, Object> updateData) throws IOException {
		UpdateByQueryRequest req = new UpdateByQueryRequest(index);
		BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
		queryBuilder.must(QueryBuilders.typeQuery(type));
		for (Entry<String, Object> entry : filterData.entrySet()) {
			queryBuilder.must(QueryBuilders.termQuery(entry.getKey(), entry.getValue()));
		}
		req.setQuery(queryBuilder);

		int upFlag = 0;
		int upSize = updateData.size();
		StringBuilder sb = new StringBuilder();
		for (Map.Entry<String, Object> entry : updateData.entrySet()) {
			upFlag++;
			String key = entry.getKey();
			sb.append("ctx._source." + key + " = params." + key);
			if (upFlag < upSize) {
				sb.append(";");
			}
		}
		req.setScript(new Script(ScriptType.INLINE, "painless", sb.toString(), updateData)).setRefresh(true)
				.indices(index);
		BulkByScrollResponse res = client.updateByQuery(req, RequestOptions.DEFAULT);
		return res.getUpdated();

	}

	public static long deleteByQuery(RestHighLevelClient client, String index, String type,
			Map<String, Object> filterData) throws IOException {
		DeleteByQueryRequest req = new DeleteByQueryRequest(index);
		BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
		queryBuilder.must(QueryBuilders.typeQuery(type));
		for (Entry<String, Object> entry : filterData.entrySet()) {
			queryBuilder.must(QueryBuilders.termQuery(entry.getKey(), entry.getValue()));
		}
		req.setQuery(queryBuilder).setRefresh(true).indices(index);
		BulkByScrollResponse res = client.deleteByQuery(req, RequestOptions.DEFAULT);
		return res.getDeleted();
	}

public static void refreshIndex(RestHighLevelClient client, String index) throws IOException {
		RefreshRequest req = new RefreshRequest(index);
		RefreshResponse refreshResponse = client.indices().refresh(req, RequestOptions.DEFAULT);
		if (refreshResponse.getFailedShards() > 0) {
			throw new KDException(BosErrorCode.fulltextException, "refresh failed:" + refreshResponse.toString());
		}
	}
	
	//for test
	public static List<Map<String, Object>> queryAll(RestHighLevelClient client, String indexName, String typeName)
			throws IOException {
		SearchRequest req = new SearchRequest(indexName);
		SearchSourceBuilder builder = new SearchSourceBuilder();
		req.types(typeName);
		builder.query(QueryBuilders.matchAllQuery());
		req.source(builder);
		SearchResponse res = client.search(req, RequestOptions.DEFAULT);
		SearchHit[] hits = res.getHits().getHits();

		List<Map<String, Object>> result = new ArrayList<>();
		for (SearchHit hit : hits) {
			result.add(hit.getSourceAsMap());
		}
		return result;

	}

public static Map<String, List<Map<String, Object>>> query(RestHighLevelClient client, String index, String[] types,
			String[] selectFields, String mustOrShould, QueryBuilder[] queryBuilders,
			Map<String, SortOrder> sortFieldMap, boolean score, int start, int size) throws IOException {
		BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
		BoolQueryBuilder filterQueryBuilder = QueryBuilders.boolQuery();
		mustOrShould(mustOrShould, queryBuilders, filterQueryBuilder);

		isNeedScore(score, queryBuilder, filterQueryBuilder);
		SearchRequest req = new SearchRequest(index);
		req.types(types);
		SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
		sourceBuilder.query(queryBuilder).from(start).size(size).fetchSource(selectFields, null);

		sort(sortFieldMap, sourceBuilder);
		req.source(sourceBuilder);
		SearchResponse searchResponse = client.search(req, RequestOptions.DEFAULT);
		SearchHit[] hits = searchResponse.getHits().getHits();
		Map<String, List<Map<String, Object>>> resultMap = new HashMap<>();
		for (String type : types) {
			resultMap.put(type, new ArrayList<>());
		}
		for (SearchHit hit : hits) {
			String type = hit.getType();
			resultMap.get(type).add(hit.getSourceAsMap());
		}
		return resultMap;
	}

	private static void sort(Map<String, SortOrder> sortFieldMap, SearchSourceBuilder builder) {
		if (sortFieldMap != null && !sortFieldMap.isEmpty()) {
			for (Entry<String, SortOrder> entry : sortFieldMap.entrySet()) {
				builder.sort(entry.getKey(), entry.getValue());
			}
		}
	}

	private static void isNeedScore(boolean score, BoolQueryBuilder queryBuilder, BoolQueryBuilder filterQb) {
		if (score) {
			queryBuilder.must(filterQb);
		} else {
			queryBuilder.filter(filterQb);
		}
	}

private static void mustOrShould(String mustOrShould, QueryBuilder[] queryBuilders, BoolQueryBuilder filterQb) {
		String must = "must";
		String should = "should";
		if ((!must.equals(mustOrShould)) && (!should.equals(mustOrShould))) {
			throw new KDException(BosErrorCode.fulltextException, "the relation ("+mustOrShould+") does not support.");
		}
		if (must.equals(mustOrShould)) {
			for (QueryBuilder qb : queryBuilders) {
				filterQb.must(qb);
			}
		} else {
			for (QueryBuilder qb : queryBuilders) {
				filterQb.should(qb);
			}
		}

	}

	public static Map<String, List<Map<String, Object>>> queryHighLight(RestHighLevelClient client, String index,
			String[] types, String[] selectFields, String mustOrShould, QueryBuilder[] queryBuilders,
			Map<String, SortOrder> sortFieldMap, boolean score, String preTag, String postTag, int start, int size)
			throws IOException {
		BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
		BoolQueryBuilder filterQueryBuilder = QueryBuilders.boolQuery();
		mustOrShould(mustOrShould, queryBuilders, filterQueryBuilder);

		isNeedScore(score, queryBuilder, filterQueryBuilder);
		SearchRequest req = new SearchRequest(index);
		req.types(types);
		SearchSourceBuilder ssb = new SearchSourceBuilder();
		ssb.query(queryBuilder).from(start).size(size).fetchSource(selectFields, null);
		sort(sortFieldMap, ssb);
		// 高亮 match生效
		List<String> highLightFields = highLighMatch(selectFields, preTag, postTag, ssb);
		req.source(ssb);
		SearchResponse searchResponse = client.search(req, RequestOptions.DEFAULT);
		return queryResult(types, highLightFields, searchResponse);

	}

private static Map<String, List<Map<String, Object>>> queryResult(String[] types, List<String> highLightFields,
			SearchResponse searchResponse) {
		SearchHit[] hits = searchResponse.getHits().getHits();
		Map<String, List<Map<String, Object>>> resultMap = new HashMap<>();
		for (String type : types) {
			resultMap.put(type, new ArrayList<>());
		}
		for (SearchHit hit : hits) {
			String type = hit.getType();
			Map<String, Object> source = hit.getSourceAsMap();
			Map<String, HighlightField> highlightFields = hit.getHighlightFields();
			for (String field : highLightFields) {
				HighlightField nameField = highlightFields.get(field);
				if (nameField != null) {
					Text[] fragments = nameField.fragments();
					String nameTmp = "";
					for (Text text : fragments) {
						nameTmp += text;
					}
					source.put(field, nameTmp);
				}
			}
			resultMap.get(type).add(source);
		}
		return resultMap;
	}

	private static List<String> highLighMatch(String[] selectFields, String preTag, String postTag,
			SearchSourceBuilder ssb) {
		List<String> highLightFields = new ArrayList<>();
		for (String field : selectFields) {
			if (!(field.equalsIgnoreCase(ConstantKeys.ENTITY_NAME_FIELD + "_String")
					&& field.equalsIgnoreCase(ConstantKeys.ENTITY_PKID_FIELD + "_String"))) {
				highLightFields.add(field);
			}
		}

		HighlightBuilder hb = new HighlightBuilder();
		hb.preTags(preTag).postTags(postTag);
		for (String field : highLightFields) {
			hb.field(field);
		}
		ssb.highlighter(hb);
		return highLightFields;
	}

public static long getCount(RestHighLevelClient client, String index, String[] types, String mustOrShould,
			QueryBuilder[] queryBuilders) throws IOException {
		BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();

		BoolQueryBuilder filterQueryBuilder = QueryBuilders.boolQuery();
		mustOrShould(mustOrShould, queryBuilders, filterQueryBuilder);
		queryBuilder.filter(filterQueryBuilder);
		SearchRequest req = new SearchRequest(index);
		req.types(types);
		SearchSourceBuilder builder = new SearchSourceBuilder();
		builder.query(queryBuilder).size(0);
		req.source(builder);
		SearchResponse searchResponse = client.search(req, RequestOptions.DEFAULT);
		return searchResponse.getHits().getTotalHits();
	}

}