es工具类

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.piesat.engine.common.base.exception.BizException;
import com.piesat.engine.hjj.data.api.common.PagerRes;
import com.piesat.engine.hjj.data.server.common.ESErrorConstant;
import com.piesat.engine.hjj.data.server.config.ESBeanConfig;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.elasticsearch.action.admin.indices.alias.IndicesAliasesRequest;
import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.action.support.master.AcknowledgedResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.client.indices.CreateIndexResponse;
import org.elasticsearch.client.indices.GetIndexRequest;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilder;
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.script.mustache.SearchTemplateRequest;
import org.elasticsearch.script.mustache.SearchTemplateResponse;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;

@Component
@Slf4j
public class ESUtil {

    private RestHighLevelClient restHighLevelClient;

    @Autowired
    private ApplicationContext applicationContext;

    @PostConstruct
    private void init() {
        try {
            restHighLevelClient = applicationContext.getBean(ESBeanConfig.class).getObject();
        } catch (Exception exp) {
            exp.printStackTrace();
            log.error(exp.getMessage());
        }
    }

    /**
     * 判断索引是否存在
     *
     * @param indexName
     * @return
     */
    public boolean isIndexExists(String indexName) {
        boolean exists = false;
        try {
            GetIndexRequest getIndexRequest = new GetIndexRequest(indexName);
            getIndexRequest.humanReadable(true);
            exists = restHighLevelClient.indices().exists(getIndexRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            log.error("判断索引是否存在发生异常", e);
        }
        return exists;
    }

    /**
     * 创建索引
     *
     * @param indexName   索引名称
     * @param mappingJson 索引映射
     */
    public void createIndex(String indexName, String mappingJson) {
        if (isIndexExists(indexName)) {
            log.info("{}索引已经存在，不能创建", indexName);
            throw new BizException(ESErrorConstant.ES_INDEX_EXISTS);
        }
        CreateIndexRequest request = new CreateIndexRequest(indexName);
        try {
            JSONObject jsonObject = JSON.parseObject(mappingJson);
            String mappings = jsonObject.getString("mappings");
            String settings = jsonObject.getString("settings");
            if (StringUtils.isNotBlank(mappings)) {
                request.mapping(mappings, XContentType.JSON);
            }
            if (StringUtils.isNotBlank(settings)) {
                request.settings(settings, XContentType.JSON);
            }
            CreateIndexResponse createIndexResponse = restHighLevelClient.indices().create(request, RequestOptions.DEFAULT);
            boolean acknowledged = createIndexResponse.isAcknowledged();
            if (!acknowledged) {
                log.error("创建索引失败，indexName={}", indexName);
                throw new BizException(ESErrorConstant.ES_INDEX_CREATE_FAILURE);
            }
        } catch (Exception e) {
            log.error("创建索引异常，indexName={}", indexName, e);
            throw new BizException(ESErrorConstant.ES_INDEX_CREATE_FAILURE);
        }

    }

    /**
     * 删除索引
     *
     * @param indexName
     */
    public void delIndex(String indexName) {
        boolean acknowledged;
        try {
            DeleteIndexRequest delRequest = new DeleteIndexRequest(indexName);
            AcknowledgedResponse delResponse = restHighLevelClient.indices().delete(delRequest, RequestOptions.DEFAULT);
            acknowledged = delResponse.isAcknowledged();
        } catch (Exception e) {
            e.printStackTrace();
            throw new BizException(ESErrorConstant.ES_INDEX_DEL_FAILURE);
        }
        if (!acknowledged) {
            log.error("删除索引失败，indexName={}", indexName);
            throw new BizException(ESErrorConstant.ES_INDEX_DEL_FAILURE);
        }
    }

    /**
     * 写入文档数据
     *
     * @param indexName
     */
    public void insertDocument(String indexName, Map<String, Object> data) {
        try {
            IndexRequest indexRequest = new IndexRequest(indexName);
            indexRequest.source(JSON.toJSONString(data), XContentType.JSON);
            restHighLevelClient.index(indexRequest, RequestOptions.DEFAULT);
        } catch (Exception e) {
            log.error("插入ES Document异常，data:{}", JSON.toJSONString(data), e);
            throw new BizException(ESErrorConstant.ES_DOC_ADD_FAILURE);
        }
    }

    /**
     * 更新文档
     *
     * @param indexName
     * @param docId
     * @param data
     */
    public void updateDocument(String indexName, String docId, Map<String, Object> data) {
        UpdateRequest updateRequest = null;
        updateRequest = new UpdateRequest(indexName, docId);
        updateRequest.doc(data);
        try {
            restHighLevelClient.update(updateRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            log.error("更新ES Document异常，data:{}", JSON.toJSONString(data), e);
            throw new BizException(ESErrorConstant.ES_DOC_UPDATE_FAILURE);
        }
    }

    /**
     * 更新文档
     *
     * @param indexName
     * @param query
     * @param data
     */
    public void updateDocumentByQuery(String indexName, QueryBuilder query, Map<String, Object> data) {
        UpdateByQueryRequest updateByQueryRequest = new UpdateByQueryRequest(indexName);
        updateByQueryRequest.setQuery(query);
        StringBuilder script = new StringBuilder();
        Set<String> keys = data.keySet();
        for (String key : keys) {
            String appendValue = "";
            Object value = data.get(key);
            if (value instanceof Number) {
                appendValue = value.toString();
            } else if (value instanceof String) {
                appendValue = "'" + value.toString() + "'";
            } else if (value instanceof List) {
                appendValue = JSONObject.toJSONString(value);
            } else {
                appendValue = value.toString();
            }
            script.append("ctx._source.").append(key).append("=").append(appendValue).append(";");
        }
        updateByQueryRequest.setScript(new Script(script.toString()));

        try {
            BulkByScrollResponse res = restHighLevelClient.updateByQuery(updateByQueryRequest, RequestOptions.DEFAULT);
            log.info("updateDocumentByQuery res:{}",res);
        } catch (IOException e) {
            log.error("更新ES Document异常，data:{}", JSON.toJSONString(data), e);
            throw new BizException(ESErrorConstant.ES_DOC_UPDATE_FAILURE);
        }
    }

    /**
     * 删除文档
     *
     * @param indexName
     * @param docId
     */
    public void delDocument(String indexName, String docId) {
        DeleteRequest deleteRequest = new DeleteRequest(indexName);
        deleteRequest.index(indexName);
        deleteRequest.id(docId);
        // 设置超时时间
        deleteRequest.timeout(TimeValue.timeValueMinutes(2));
        // 设置刷新策略"wait_for"
        // 保持此请求打开，直到刷新使此请求的内容可以搜索为止。此刷新策略与高索引和搜索吞吐量兼容，但它会导致请求等待响应，直到发生刷新
        deleteRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.WAIT_UNTIL);
        // 同步删除
        DeleteResponse deleteResponse = null;
        try {
            deleteResponse = restHighLevelClient.delete(deleteRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            log.error("删除ES Document异常，docId={}", docId, e);
            throw new BizException(ESErrorConstant.ES_DOC_DEL_FAILURE);
        }
        int status = deleteResponse.status().getStatus();
        if (status != 200) {
            log.error("删除文档失败，indexName={},docId={}", indexName, docId);
            throw new BizException(ESErrorConstant.ES_DOC_DEL_FAILURE);
        }
    }

    /**
     * 通过条件删除
     *
     * @param indexName
     * @param query
     */
    public void delDocumentByQuery(String indexName, QueryBuilder query) {
        DeleteByQueryRequest deleteRequest = new DeleteByQueryRequest(indexName);
        deleteRequest.setQuery(query);
        BulkByScrollResponse deleteResponse = null;
        try {
            deleteResponse = restHighLevelClient.deleteByQuery(deleteRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            log.error("删除ES Document异常，query = {}", query, e);
            throw new BizException(ESErrorConstant.ES_DOC_DEL_FAILURE);
        }
    }

    /**
     * 通过查询模板查询数据
     *
     * @param indexAlias
     * @param templateJson
     * @param paramMap
     * @return
     * @throws IOException
     */
    public PagerRes<String> queryDataByTemplate(String indexAlias, String templateJson, Map<String, Object> paramMap) throws IOException {
        SearchTemplateRequest request = new SearchTemplateRequest();
        request.setRequest(new SearchRequest(indexAlias));
        request.setScriptType(ScriptType.INLINE);
        request.setScript(templateJson);
        request.setScriptParams(paramMap);

        List<String> resList = new ArrayList<>();
        SearchTemplateResponse searchTemplateResponse = restHighLevelClient.searchTemplate(request, RequestOptions.DEFAULT);
        SearchHit[] hits = searchTemplateResponse.getResponse().getHits().getHits();
        for (SearchHit hit : hits) {
            String jsonData = hit.getSourceAsString();
            resList.add(jsonData);
        }
        return PagerRes.<String>builder().data(resList).total(searchTemplateResponse.getResponse().getHits().getTotalHits().value).build();
    }

    /**
     * 设置索引别名
     *
     * @param oldIndexName
     * @param newIndexName
     * @param alias
     */
    public void setIndexAlias(String oldIndexName, String newIndexName, String alias) throws IOException {
        IndicesAliasesRequest request = new IndicesAliasesRequest();
        if (StringUtils.isNotBlank(oldIndexName)) {
            IndicesAliasesRequest.AliasActions oldAliasAction =
                    new IndicesAliasesRequest.AliasActions(IndicesAliasesRequest.AliasActions.Type.REMOVE)
                            .index(oldIndexName)
                            .alias(alias);
            request.addAliasAction(oldAliasAction);
        }
        if (StringUtils.isNotBlank(newIndexName)) {
            IndicesAliasesRequest.AliasActions newAliasAction =
                    new IndicesAliasesRequest.AliasActions(IndicesAliasesRequest.AliasActions.Type.ADD)
                            .index(newIndexName)
                            .alias(alias);
            request.addAliasAction(newAliasAction);
        }
        AcknowledgedResponse indicesAliasesResponse = restHighLevelClient.indices().updateAliases(request, RequestOptions.DEFAULT);
        boolean acknowledged = indicesAliasesResponse.isAcknowledged();
        if (!acknowledged) {
            log.error("设置索引别名失败，oldIndexName={},newIndexName={},alias={}", oldIndexName, newIndexName, alias);
            throw new BizException(ESErrorConstant.ES_INDEX_SETALIAS_FAILURE);
        }
    }

    public SearchHits commonQueryByFilter(String indexName, SearchSourceBuilder searchSourceBuilder) {
        try {
            SearchRequest searchRequest = new SearchRequest(indexName).source(searchSourceBuilder);
            SearchResponse searchResponse = restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            return searchResponse.getHits();
        } catch (Exception e) {
            e.printStackTrace();
            throw new BizException(ESErrorConstant.QUERY_ERROR);
        }
    }
}