Skip to content

Commit

Permalink
Rewrite these search calls
Browse files Browse the repository at this point in the history
  • Loading branch information
joegallo committed Nov 1, 2023
1 parent 3e26637 commit e5f5d41
Showing 1 changed file with 38 additions and 63 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -11,21 +11,20 @@
import org.elasticsearch.Version;
import org.elasticsearch.action.admin.cluster.repositories.put.PutRepositoryRequest;
import org.elasticsearch.action.admin.cluster.snapshots.restore.RestoreSnapshotRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Request;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.ResponseException;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.cluster.routing.Murmur3HashFunction;
import org.elasticsearch.common.Strings;
import org.elasticsearch.common.document.DocumentField;
import org.elasticsearch.common.settings.SecureString;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.util.concurrent.ThreadContext;
import org.elasticsearch.core.Booleans;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.core.PathUtils;
import org.elasticsearch.index.IndexVersion;
import org.elasticsearch.index.query.QueryBuilders;
Expand All @@ -42,7 +41,6 @@

import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
Expand Down Expand Up @@ -74,12 +72,6 @@ protected Settings restClientSettings() {
return Settings.builder().put(ThreadContext.PREFIX + ".Authorization", token).build();
}

@SuppressWarnings("removal")
protected static RestHighLevelClient highLevelClient(RestClient client) {
return new RestHighLevelClient(client, ignore -> {}, Collections.emptyList()) {
};
}

public void testOldRepoAccess() throws IOException {
runTest(false);
}
Expand Down Expand Up @@ -112,10 +104,7 @@ public void runTest(boolean sourceOnlyRepository) throws IOException {
int numDocs = 10;
int extraDocs = 1;
final Set<String> expectedIds = new HashSet<>();
try (
RestHighLevelClient client = highLevelClient(adminClient());
RestClient oldEs = RestClient.builder(new HttpHost("127.0.0.1", oldEsPort)).build()
) {
try (RestClient oldEs = RestClient.builder(new HttpHost("127.0.0.1", oldEsPort)).build()) {
if (afterRestart == false) {
beforeRestart(
sourceOnlyRepository,
Expand All @@ -125,7 +114,6 @@ public void runTest(boolean sourceOnlyRepository) throws IOException {
numDocs,
extraDocs,
expectedIds,
client,
oldEs,
indexName
);
Expand All @@ -150,7 +138,6 @@ private void beforeRestart(
int numDocs,
int extraDocs,
Set<String> expectedIds,
RestHighLevelClient client,
RestClient oldEs,
String indexName
) throws IOException {
Expand Down Expand Up @@ -261,35 +248,15 @@ private void beforeRestart(
assertThat(getResp.evaluate("snapshots.0.stats.total.file_count"), greaterThan(0));

// restore / mount and check whether searches work
restoreMountAndVerify(
numDocs,
expectedIds,
client,
numberOfShards,
sourceOnlyRepository,
oldVersion,
indexName,
repoName,
snapshotName
);
restoreMountAndVerify(numDocs, expectedIds, numberOfShards, sourceOnlyRepository, oldVersion, indexName, repoName, snapshotName);

// close indices
closeIndex(client(), "restored_" + indexName);
closeIndex(client(), "mounted_full_copy_" + indexName);
closeIndex(client(), "mounted_shared_cache_" + indexName);

// restore / mount again
restoreMountAndVerify(
numDocs,
expectedIds,
client,
numberOfShards,
sourceOnlyRepository,
oldVersion,
indexName,
repoName,
snapshotName
);
restoreMountAndVerify(numDocs, expectedIds, numberOfShards, sourceOnlyRepository, oldVersion, indexName, repoName, snapshotName);
}

private String getType(Version oldVersion, String id) {
Expand All @@ -304,7 +271,6 @@ private static String sourceForDoc(int i) {
private void restoreMountAndVerify(
int numDocs,
Set<String> expectedIds,
RestHighLevelClient client,
int numberOfShards,
boolean sourceOnlyRepository,
Version oldVersion,
Expand Down Expand Up @@ -357,7 +323,7 @@ private void restoreMountAndVerify(
}

// run a search against the index
assertDocs("restored_" + indexName, numDocs, expectedIds, client, sourceOnlyRepository, oldVersion, numberOfShards);
assertDocs("restored_" + indexName, numDocs, expectedIds, sourceOnlyRepository, oldVersion, numberOfShards);

// mount as full copy searchable snapshot
Request mountRequest = new Request("POST", "/_snapshot/" + repoName + "/" + snapshotName + "/_mount");
Expand All @@ -377,7 +343,7 @@ private void restoreMountAndVerify(
ensureGreen("mounted_full_copy_" + indexName);

// run a search against the index
assertDocs("mounted_full_copy_" + indexName, numDocs, expectedIds, client, sourceOnlyRepository, oldVersion, numberOfShards);
assertDocs("mounted_full_copy_" + indexName, numDocs, expectedIds, sourceOnlyRepository, oldVersion, numberOfShards);

// mount as shared cache searchable snapshot
mountRequest = new Request("POST", "/_snapshot/" + repoName + "/" + snapshotName + "/_mount");
Expand All @@ -390,15 +356,14 @@ private void restoreMountAndVerify(
assertEquals(numberOfShards, (int) mountResponse.evaluate("snapshot.shards.successful"));

// run a search against the index
assertDocs("mounted_shared_cache_" + indexName, numDocs, expectedIds, client, sourceOnlyRepository, oldVersion, numberOfShards);
assertDocs("mounted_shared_cache_" + indexName, numDocs, expectedIds, sourceOnlyRepository, oldVersion, numberOfShards);
}

@SuppressWarnings("removal")
private void assertDocs(
String index,
int numDocs,
Set<String> expectedIds,
RestHighLevelClient client,
boolean sourceOnlyRepository,
Version oldVersion,
int numberOfShards
Expand All @@ -409,8 +374,10 @@ private void assertDocs(
.build();
RequestOptions randomRequestOptions = randomBoolean() ? RequestOptions.DEFAULT : v7RequestOptions;

SearchResponse searchResponse;

// run a search against the index
SearchResponse searchResponse = client.search(new SearchRequest(index), randomRequestOptions);
searchResponse = search(index, null, randomRequestOptions);
logger.info(searchResponse);
// check hit count
assertEquals(numDocs, searchResponse.getHits().getTotalHits().value);
Expand All @@ -428,12 +395,11 @@ private void assertDocs(
String id = randomFrom(expectedIds);
int num = getIdAsNumeric(id);
// run a search using runtime fields against the index
searchResponse = client.search(
new SearchRequest(index).source(
SearchSourceBuilder.searchSource()
.query(QueryBuilders.matchQuery("val", num))
.runtimeMappings(Map.of("val", Map.of("type", "long")))
),
searchResponse = search(
index,
SearchSourceBuilder.searchSource()
.query(QueryBuilders.matchQuery("val", num))
.runtimeMappings(Map.of("val", Map.of("type", "long"))),
randomRequestOptions
);
logger.info(searchResponse);
Expand All @@ -443,12 +409,11 @@ private void assertDocs(

if (sourceOnlyRepository == false) {
// search using reverse sort on val
searchResponse = client.search(
new SearchRequest(index).source(
SearchSourceBuilder.searchSource()
.query(QueryBuilders.matchAllQuery())
.sort(SortBuilders.fieldSort("val").order(SortOrder.DESC))
),
searchResponse = search(
index,
SearchSourceBuilder.searchSource()
.query(QueryBuilders.matchAllQuery())
.sort(SortBuilders.fieldSort("val").order(SortOrder.DESC)),
randomRequestOptions
);
logger.info(searchResponse);
Expand All @@ -459,8 +424,9 @@ private void assertDocs(
);

// look up postings
searchResponse = client.search(
new SearchRequest(index).source(SearchSourceBuilder.searchSource().query(QueryBuilders.matchQuery("test", "test" + num))),
searchResponse = search(
index,
SearchSourceBuilder.searchSource().query(QueryBuilders.matchQuery("test", "test" + num)),
randomRequestOptions
);
logger.info(searchResponse);
Expand All @@ -471,8 +437,9 @@ private void assertDocs(
// search on _type and check that results contain _type information
String randomType = getType(oldVersion, randomFrom(expectedIds));
long typeCount = expectedIds.stream().filter(idd -> getType(oldVersion, idd).equals(randomType)).count();
searchResponse = client.search(
new SearchRequest(index).source(SearchSourceBuilder.searchSource().query(QueryBuilders.termQuery("_type", randomType))),
searchResponse = search(
index,
SearchSourceBuilder.searchSource().query(QueryBuilders.termQuery("_type", randomType)),
randomRequestOptions
);
logger.info(searchResponse);
Expand All @@ -492,10 +459,9 @@ private void assertDocs(
);

// check that shards are skipped based on non-matching date
searchResponse = client.search(
new SearchRequest(index).source(
SearchSourceBuilder.searchSource().query(QueryBuilders.rangeQuery("create_date").from("2020-02-01"))
),
searchResponse = search(
index,
SearchSourceBuilder.searchSource().query(QueryBuilders.rangeQuery("create_date").from("2020-02-01")),
randomRequestOptions
);
logger.info(searchResponse);
Expand All @@ -506,6 +472,15 @@ private void assertDocs(
}
}

private static SearchResponse search(String index, @Nullable SearchSourceBuilder builder, RequestOptions options) throws IOException {
Request request = new Request("POST", "/" + index + "/_search");
if (builder != null) {
request.setJsonEntity(builder.toString());
}
request.setOptions(options);
return SearchResponse.fromXContent(responseAsParser(client().performRequest(request)));
}

private int getIdAsNumeric(String id) {
return Integer.parseInt(id.substring("testdoc".length()));
}
Expand Down

0 comments on commit e5f5d41

Please sign in to comment.