package write;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.http.HttpHost;
import org.elasticsearch.action.bulk.BackoffPolicy;
import org.elasticsearch.action.bulk.BulkProcessor;
import org.elasticsearch.action.bulk.BulkRequest;
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.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.client.indices.CreateIndexResponse;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.unit.ByteSizeUnit;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.MatchAllQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.aggregations.bucket.sampler.Sampler;
import org.elasticsearch.search.aggregations.bucket.sampler.SamplerAggregationBuilder;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.junit.jupiter.api.Test;
import java.io.IOException;
import java.util.concurrent.TimeUnit;
@Slf4j
public class IndexTest {
RestHighLevelClient newClient() {
return new RestHighLevelClient(RestClient.builder(new HttpHost("localhost", 9200, "http")));
}
BulkProcessor newBulkProcessor(RestHighLevelClient client) {
BulkProcessor.Listener listener = new BulkProcessor.Listener() {
@Override
public void beforeBulk(long executionId, BulkRequest request) {
}
@Override
public void afterBulk(long executionId, BulkRequest request,
BulkResponse response) {
if (response.hasFailures()) {
log.warn("Bulk [{}] executed with failures", executionId);
} else {
log.info("Bulk [{}] completed in {} milliseconds",
executionId, response.getTook().getMillis());
}
}
@Override
public void afterBulk(long executionId, BulkRequest request,
Throwable failure) {
log.error("Failed to execute bulk", failure);
}
};
BulkProcessor.Builder builder = BulkProcessor.builder((request, bulkListener) -> client.bulkAsync(request, RequestOptions.DEFAULT, bulkListener),
listener);
builder.setBulkActions(2000);
builder.setBulkSize(new ByteSizeValue(1L, ByteSizeUnit.MB));
builder.setConcurrentRequests(5);
builder.setFlushInterval(TimeValue.timeValueSeconds(10L));
builder.setBackoffPolicy(BackoffPolicy
.constantBackoff(TimeValue.timeValueSeconds(1L), 3));
return builder.build();
}
@Test
void testWrite() throws InterruptedException, IOException {
String index = "test1";
RestHighLevelClient client = newClient();
BulkProcessor bulkProcessor = newBulkProcessor(client);
int count = 100;
for (int i = 0; i < 100000000; i++) {
String randomAlphanumeric = RandomStringUtils.randomAlphanumeric(200);
IndexRequest indexRequest = new IndexRequest(index).source(XContentType.JSON, "content", randomAlphanumeric, "time", System.currentTimeMillis());
bulkProcessor.add(indexRequest);
for (int j = 0; j < count; j++) {
bulkProcessor.add(indexRequest);
}
count--;
}
bulkProcessor.awaitClose(30L, TimeUnit.SECONDS);
}
@Test
void testCreateIndex() throws IOException {
RestHighLevelClient client = newClient();
CreateIndexRequest request = new CreateIndexRequest("test1");
request.settings(Settings.builder()
.put("index.number_of_shards", 12)
.put("index.number_of_replicas", 0)
);
CreateIndexResponse createIndexResponse = client.indices().create(request, RequestOptions.DEFAULT);
log.info("create index {}", createIndexResponse.isAcknowledged());
client.close();
}
@Test
void testCount() throws IOException {
RestHighLevelClient client = newClient();
CountRequest countRequest = new CountRequest("test6");
countRequest.query(new MatchAllQueryBuilder());
CountResponse count = client.count(countRequest, RequestOptions.DEFAULT);
System.out.println(count.getCount());
client.close();
}
@Test
void testSample() throws IOException {
RestHighLevelClient client = newClient();
RangeQueryBuilder rangeQueryBuilder = QueryBuilders.rangeQuery("time").gte(1658312038698L).lte(1658312038698L);
TermsAggregationBuilder termsAggregationBuilder = new TermsAggregationBuilder("contents").field("content.keyword").size(100);
SamplerAggregationBuilder aggBuilder = new SamplerAggregationBuilder("sampler").shardSize(100000)
.subAggregation(termsAggregationBuilder);
SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
searchSourceBuilder.query(QueryBuilders.boolQuery().filter(rangeQueryBuilder));
searchSourceBuilder.aggregation(aggBuilder);
searchSourceBuilder.size(0);
searchSourceBuilder.timeout(TimeValue.timeValueMillis(100000));
SearchRequest searchRequest = new SearchRequest();
log.info("sourceBuilder: {}", searchSourceBuilder);
searchRequest.source(searchSourceBuilder);
searchRequest.indices("test6");
SearchResponse response = client.search(searchRequest, RequestOptions.DEFAULT);
System.out.println(response.toString());
Sampler sampler = (Sampler) response.getAggregations().asMap().get("sampler");
Terms contents = (Terms) sampler.getAggregations().get("contents");
for (Terms.Bucket bucket : contents.getBuckets()) {
System.out.println(bucket.getKey());
}
client.close();
}
@Test
void testTermAggregation() throws IOException {
RestHighLevelClient client = newClient();
RangeQueryBuilder rangeQueryBuilder = QueryBuilders.rangeQuery("time").gte(1658312030698L).lte(1658312038698L);
TermsAggregationBuilder termsAggregationBuilder = new TermsAggregationBuilder("contents").field("content.keyword").size(100);
SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
searchSourceBuilder.query(new MatchAllQueryBuilder());
searchSourceBuilder.aggregation(termsAggregationBuilder);
searchSourceBuilder.timeout(TimeValue.timeValueMillis(10000));
SearchRequest searchRequest = new SearchRequest();
log.info("sourceBuilder: {}", searchSourceBuilder);
searchRequest.source(searchSourceBuilder);
searchRequest.indices("test6");
SearchResponse response = client.search(searchRequest, RequestOptions.DEFAULT);
System.out.println(response.toString());
Terms contents = (Terms) response.getAggregations().get("contents");
for (Terms.Bucket bucket : contents.getBuckets()) {
System.out.println(bucket.getKey());
}
client.close();
}
@Test
void testSearch() throws IOException {
RestHighLevelClient client = newClient();
SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
RangeQueryBuilder rangeQueryBuilder = QueryBuilders.rangeQuery("time").gte(1658312038698L).lte(1658312038698L);
searchSourceBuilder.query(rangeQueryBuilder);
searchSourceBuilder.timeout(TimeValue.timeValueMillis(10000));
SearchRequest searchRequest = new SearchRequest();
log.info("sourceBuilder: {}", searchSourceBuilder);
searchRequest.source(searchSourceBuilder);
searchRequest.indices("test6");
SearchResponse response = client.search(searchRequest, RequestOptions.DEFAULT);
System.out.println(response.toString());
for (SearchHit hit : response.getHits().getHits()) {
log.info("doc : {}", hit.getSourceAsString());
}
client.close();
}
}
