org.apache.pulsar.io.elasticsearch.client.opensearch.OpenSearchHighLevelRestClient Maven / Gradle / Ivy
/*
* 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.pulsar.io.elasticsearch.client.opensearch;
import com.google.common.annotations.VisibleForTesting;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.apache.http.HttpHost;
import org.apache.pulsar.functions.api.Record;
import org.apache.pulsar.io.elasticsearch.ElasticSearchConfig;
import org.apache.pulsar.io.elasticsearch.RandomExponentialRetry;
import org.apache.pulsar.io.elasticsearch.client.BulkProcessor;
import org.apache.pulsar.io.elasticsearch.client.RestClient;
import org.opensearch.action.DocWriteResponse;
import org.opensearch.action.admin.indices.delete.DeleteIndexRequest;
import org.opensearch.action.admin.indices.refresh.RefreshRequest;
import org.opensearch.action.bulk.BulkRequest;
import org.opensearch.action.bulk.BulkResponse;
import org.opensearch.action.delete.DeleteRequest;
import org.opensearch.action.delete.DeleteResponse;
import org.opensearch.action.index.IndexRequest;
import org.opensearch.action.index.IndexResponse;
import org.opensearch.action.search.SearchRequest;
import org.opensearch.action.search.SearchResponse;
import org.opensearch.client.RequestOptions;
import org.opensearch.client.Requests;
import org.opensearch.client.RestClientBuilder;
import org.opensearch.client.RestHighLevelClient;
import org.opensearch.client.indices.CreateIndexRequest;
import org.opensearch.client.indices.CreateIndexResponse;
import org.opensearch.client.indices.GetIndexRequest;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.unit.TimeValue;
import org.opensearch.common.xcontent.XContentType;
import org.opensearch.core.common.Strings;
import org.opensearch.core.common.unit.ByteSizeUnit;
import org.opensearch.core.common.unit.ByteSizeValue;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryBuilders;
import org.opensearch.search.builder.SearchSourceBuilder;
@Slf4j
public class OpenSearchHighLevelRestClient extends RestClient implements BulkProcessor {
private interface DocWriteRequestWithPulsarRecord {
Record getPulsarRecord();
}
private static class IndexRequestWithPulsarRecord extends IndexRequest implements DocWriteRequestWithPulsarRecord {
private Record pulsarRecord;
public IndexRequestWithPulsarRecord(String index, Record pulsarRecord) {
super(index);
this.pulsarRecord = pulsarRecord;
}
@Override
public Record getPulsarRecord() {
return pulsarRecord;
}
}
private static class DeleteRequestWithPulsarRecord
extends DeleteRequest
implements DocWriteRequestWithPulsarRecord {
private Record pulsarRecord;
public DeleteRequestWithPulsarRecord(String index, Record pulsarRecord) {
super(index);
this.pulsarRecord = pulsarRecord;
}
@Override
public Record getPulsarRecord() {
return pulsarRecord;
}
}
private RestHighLevelClient client;
private org.opensearch.action.bulk.BulkProcessor internalBulkProcessor;
public OpenSearchHighLevelRestClient(ElasticSearchConfig elasticSearchConfig,
BulkProcessor.Listener bulkProcessorListener) {
super(elasticSearchConfig, bulkProcessorListener);
log.info("ElasticSearch URL {}", config.getElasticSearchUrl());
final HttpHost[] httpHosts = getHttpHosts();
RestClientBuilder builder = org.opensearch.client.RestClient.builder(httpHosts)
.setRequestConfigCallback(builder1 -> builder1
.setContentCompressionEnabled(config.isCompressionEnabled())
.setConnectionRequestTimeout(config.getConnectionRequestTimeoutInMs())
.setConnectTimeout(config.getConnectTimeoutInMs())
.setSocketTimeout(config.getSocketTimeoutInMs()))
.setCompressionEnabled(config.isCompressionEnabled())
.setHttpClientConfigCallback(this.configCallback)
.setFailureListener(new org.opensearch.client.RestClient.FailureListener() {
@Override
public void onFailure(org.opensearch.client.Node node) {
log.warn("Node host={} failed", node.getHost());
}
});
client = new RestHighLevelClient(builder);
if (config.isBulkEnabled()) {
org.opensearch.action.bulk.BulkProcessor.Builder bulkBuilder = org.opensearch.action.bulk.BulkProcessor
.builder(
(bulkRequest, bulkResponseActionListener)
-> client.bulkAsync(bulkRequest,
RequestOptions.DEFAULT,
bulkResponseActionListener),
new org.opensearch.action.bulk.BulkProcessor.Listener() {
private List
convertBulkRequest(BulkRequest bulkRequest) {
return bulkRequest.requests().stream().map(docWriteRequest -> {
final Record pulsarRecord;
if (docWriteRequest instanceof DocWriteRequestWithPulsarRecord) {
DocWriteRequestWithPulsarRecord requestWithId =
(DocWriteRequestWithPulsarRecord) docWriteRequest;
pulsarRecord = requestWithId.getPulsarRecord();
} else {
throw new UnsupportedOperationException("Unexpected bulk request of type: "
+ docWriteRequest.getClass());
}
return BulkProcessor.BulkOperationRequest.builder()
.pulsarRecord(pulsarRecord)
.build();
}).collect(Collectors.toList());
}
private List
convertBulkResponse(BulkResponse bulkRequest) {
return Arrays.asList(bulkRequest.getItems())
.stream()
.map(itemResponse ->
BulkProcessor.BulkOperationResult.builder()
.error(itemResponse.getFailureMessage())
.index(itemResponse.getIndex())
.documentId(itemResponse.getId())
.build())
.collect(Collectors.toList());
}
@Override
public void beforeBulk(long l, BulkRequest bulkRequest) {
}
@Override
public void afterBulk(long l, BulkRequest bulkRequest,
BulkResponse bulkResponse) {
bulkProcessorListener.afterBulk(l, convertBulkRequest(bulkRequest),
convertBulkResponse(bulkResponse));
}
@Override
public void afterBulk(long l, BulkRequest bulkRequest, Throwable throwable) {
bulkProcessorListener.afterBulk(l, convertBulkRequest(bulkRequest),
throwable);
}
}
)
.setBulkActions(config.getBulkActions())
.setBulkSize(new ByteSizeValue(config.getBulkSizeInMb(), ByteSizeUnit.MB))
.setConcurrentRequests(config.getBulkConcurrentRequests())
.setBackoffPolicy(new RandomExponentialBackoffPolicy(
new RandomExponentialRetry(elasticSearchConfig.getMaxRetryTimeInSec()),
config.getRetryBackoffInMs(),
config.getMaxRetries()
));
if (config.getBulkFlushIntervalInMs() > 0) {
bulkBuilder.setFlushInterval(new TimeValue(config.getBulkFlushIntervalInMs(), TimeUnit.MILLISECONDS));
}
this.internalBulkProcessor = bulkBuilder.build();
} else {
this.internalBulkProcessor = null;
}
}
@Override
public boolean indexExists(String index) throws IOException {
final GetIndexRequest request = new GetIndexRequest(index);
return client.indices().exists(request, RequestOptions.DEFAULT);
}
@Override
public boolean createIndex(String index) throws IOException {
final CreateIndexRequest cireq = new CreateIndexRequest(index);
cireq.settings(Settings.builder()
.put("index.number_of_shards", config.getIndexNumberOfShards())
.put("index.number_of_replicas", config.getIndexNumberOfReplicas()));
CreateIndexResponse resp = client.indices().create(cireq, RequestOptions.DEFAULT);
if (!resp.isAcknowledged() || !resp.isShardsAcknowledged()) {
throw new IOException("Unable to create index.");
}
return true;
}
@Override
public boolean deleteIndex(String index) throws IOException {
return client.indices().delete(new DeleteIndexRequest(index), RequestOptions.DEFAULT).isAcknowledged();
}
@Override
public boolean indexDocument(String index, String documentId, String documentSource) throws IOException {
IndexRequest indexRequest = Requests.indexRequest(index);
if (!Strings.isNullOrEmpty(documentId)) {
indexRequest.id(documentId);
}
indexRequest.source(documentSource, XContentType.JSON);
IndexResponse indexResponse = client.index(indexRequest, RequestOptions.DEFAULT);
if (indexResponse.getResult().equals(DocWriteResponse.Result.CREATED)
|| indexResponse.getResult().equals(DocWriteResponse.Result.UPDATED)) {
return true;
} else {
return false;
}
}
@Override
public boolean deleteDocument(String index, String documentId) throws IOException {
DeleteRequest deleteRequest = Requests.deleteRequest(index);
deleteRequest.id(documentId);
DeleteResponse deleteResponse = client.delete(deleteRequest, RequestOptions.DEFAULT);
if (log.isDebugEnabled()) {
log.debug("delete result {}", deleteResponse.getResult());
}
if (deleteResponse.getResult().equals(DocWriteResponse.Result.DELETED)
|| deleteResponse.getResult().equals(DocWriteResponse.Result.NOT_FOUND)) {
return true;
}
return false;
}
@Override
public long totalHits(String indexName) throws IOException {
return search(indexName).getHits().getTotalHits().value;
}
@Override
public long totalHits(String indexName, String query) throws IOException {
return search(indexName, query).getHits().getTotalHits().value;
}
@VisibleForTesting
public SearchResponse search(String indexName) throws IOException {
return search(indexName, "*:*");
}
@VisibleForTesting
public SearchResponse search(String indexName, String query) throws IOException {
client.indices().refresh(new RefreshRequest(indexName), RequestOptions.DEFAULT);
QueryBuilder queryBuilder;
if ("*:*".equals(query)) {
queryBuilder = QueryBuilders.matchAllQuery();
} else {
final String[] split = query.split(":");
final String name = split[0];
final String text = split[1];
queryBuilder = QueryBuilders.matchQuery(name, text);
}
return client.search(
new SearchRequest()
.indices(indexName)
.source(new SearchSourceBuilder().query(queryBuilder)) ,
RequestOptions.DEFAULT);
}
@Override
public BulkProcessor getBulkProcessor() {
return this;
}
@Override
public void appendIndexRequest(BulkProcessor.BulkIndexRequest request) throws IOException {
IndexRequest indexRequest = new IndexRequestWithPulsarRecord(request.getIndex(), request.getRecord());
if (!Strings.isNullOrEmpty(request.getDocumentId())) {
indexRequest.id(request.getDocumentId());
}
indexRequest.source(request.getDocumentSource(), XContentType.JSON);
internalBulkProcessor.add(indexRequest);
}
@Override
public void appendDeleteRequest(BulkProcessor.BulkDeleteRequest request) throws IOException {
DeleteRequest deleteRequest = new DeleteRequestWithPulsarRecord(request.getIndex(), request.getRecord());
deleteRequest.id(request.getDocumentId());
internalBulkProcessor.add(deleteRequest);
}
@Override
public void flush() {
internalBulkProcessor.flush();
}
@Override
public void closeClient() {
try {
if (internalBulkProcessor != null) {
internalBulkProcessor.awaitClose(5000, TimeUnit.MILLISECONDS);
internalBulkProcessor = null;
}
} catch (InterruptedException e) {
log.warn("Elasticsearch bulk processor close error:", e);
}
try {
if (this.client != null) {
this.client.close();
this.client = null;
}
} catch (IOException e) {
log.warn("Elasticsearch client close error:", e);
}
}
@VisibleForTesting
public void setClient(RestHighLevelClient client) {
this.client = client;
}
@VisibleForTesting
public void setInternalBulkProcessor(org.opensearch.action.bulk.BulkProcessor internalBulkProcessor) {
this.internalBulkProcessor = internalBulkProcessor;
}
@VisibleForTesting
public RestHighLevelClient getClient() {
return client;
}
@VisibleForTesting
public org.opensearch.action.bulk.BulkProcessor getInternalBulkProcessor() {
return internalBulkProcessor;
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy