Skip to content

Commit

Permalink
Ensure that the RandomizingClient doesn't create invalid request
Browse files Browse the repository at this point in the history
  • Loading branch information
jimczi committed Dec 11, 2023
1 parent c5ea8a1 commit 93063b5
Show file tree
Hide file tree
Showing 4 changed files with 39 additions and 25 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ public void testBasic() {
}
refresh("test");
String pitId = openPointInTime(new String[] { "test" }, TimeValue.timeValueMinutes(2));
assertResponse(prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pitId)), resp1 -> {
assertResponse(prepareSearch().setPointInTime(new PointInTimeBuilder(pitId)), resp1 -> {
assertThat(resp1.pointInTimeId(), equalTo(pitId));
assertHitCount(resp1, numDocs);
});
Expand All @@ -99,13 +99,13 @@ public void testBasic() {
if (randomBoolean()) {
final int delDocCount = deletedDocs;
assertNoFailuresAndResponse(
prepareSearch("test").setPreference(null).setQuery(new MatchAllQueryBuilder()),
prepareSearch("test").setQuery(new MatchAllQueryBuilder()),
resp2 -> assertHitCount(resp2, numDocs - delDocCount)
);
}
try {
assertNoFailuresAndResponse(
prepareSearch().setPreference(null).setQuery(new MatchAllQueryBuilder()).setPointInTime(new PointInTimeBuilder(pitId)),
prepareSearch().setQuery(new MatchAllQueryBuilder()).setPointInTime(new PointInTimeBuilder(pitId)),
resp3 -> {
assertHitCount(resp3, numDocs);
assertThat(resp3.pointInTimeId(), equalTo(pitId));
Expand All @@ -131,7 +131,7 @@ public void testMultipleIndices() {
String pitId = openPointInTime(new String[] { "*" }, TimeValue.timeValueMinutes(2));
try {
int moreDocs = randomIntBetween(10, 50);
assertNoFailuresAndResponse(prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertNoFailuresAndResponse(prepareSearch().setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertHitCount(resp, numDocs);
assertNotNull(resp.pointInTimeId());
assertThat(resp.pointInTimeId(), equalTo(pitId));
Expand All @@ -143,7 +143,7 @@ public void testMultipleIndices() {
refresh();
});
assertNoFailuresAndResponse(prepareSearch(), resp -> assertHitCount(resp, numDocs + moreDocs));
assertNoFailuresAndResponse(prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertNoFailuresAndResponse(prepareSearch().setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertHitCount(resp, numDocs);
assertNotNull(resp.pointInTimeId());
assertThat(resp.pointInTimeId(), equalTo(pitId));
Expand Down Expand Up @@ -212,7 +212,7 @@ public void testRelocation() throws Exception {
refresh();
String pitId = openPointInTime(new String[] { "test" }, TimeValue.timeValueMinutes(2));
try {
assertNoFailuresAndResponse(prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertNoFailuresAndResponse(prepareSearch().setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertHitCount(resp, numDocs);
assertThat(resp.pointInTimeId(), equalTo(pitId));
});
Expand All @@ -232,7 +232,7 @@ public void testRelocation() throws Exception {
}
refresh();
}
assertNoFailuresAndResponse(prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertNoFailuresAndResponse(prepareSearch().setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertHitCount(resp, numDocs);
assertThat(resp.pointInTimeId(), equalTo(pitId));
});
Expand All @@ -245,7 +245,7 @@ public void testRelocation() throws Exception {
.collect(Collectors.toSet());
assertThat(assignedNodes, everyItem(not(in(excludedNodes))));
}, 30, TimeUnit.SECONDS);
assertNoFailuresAndResponse(prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertNoFailuresAndResponse(prepareSearch().setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertHitCount(resp, numDocs);
assertThat(resp.pointInTimeId(), equalTo(pitId));
});
Expand All @@ -263,7 +263,7 @@ public void testPointInTimeNotFound() throws Exception {
}
refresh();
String pit = openPointInTime(new String[] { "index" }, TimeValue.timeValueSeconds(5));
assertNoFailuresAndResponse(prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pit)), resp1 -> {
assertNoFailuresAndResponse(prepareSearch().setPointInTime(new PointInTimeBuilder(pit)), resp1 -> {
assertHitCount(resp1, index1);
if (rarely()) {
try {
Expand All @@ -280,7 +280,7 @@ public void testPointInTimeNotFound() throws Exception {
});
SearchPhaseExecutionException e = expectThrows(
SearchPhaseExecutionException.class,
() -> prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pit)).get()
() -> prepareSearch().setPointInTime(new PointInTimeBuilder(pit)).get()
);
for (ShardSearchFailure failure : e.shardFailures()) {
assertThat(ExceptionsHelper.unwrapCause(failure.getCause()), instanceOf(SearchContextMissingException.class));
Expand All @@ -306,7 +306,7 @@ public void testIndexNotFound() {
String pit = openPointInTime(new String[] { "index-*" }, TimeValue.timeValueMinutes(2));
try {
assertNoFailuresAndResponse(
prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pit)),
prepareSearch().setPointInTime(new PointInTimeBuilder(pit)),
resp -> assertHitCount(resp, index1 + index2)
);
indicesAdmin().prepareDelete("index-1").get();
Expand All @@ -316,7 +316,7 @@ public void testIndexNotFound() {

// Allow partial search result
assertResponse(
prepareSearch().setPreference(null).setAllowPartialSearchResults(true).setPointInTime(new PointInTimeBuilder(pit)),
prepareSearch().setAllowPartialSearchResults(true).setPointInTime(new PointInTimeBuilder(pit)),
resp -> {
assertFailures(resp);
assertHitCount(resp, index2);
Expand All @@ -326,7 +326,7 @@ public void testIndexNotFound() {
// Do not allow partial search result
expectThrows(
ElasticsearchException.class,
() -> prepareSearch().setPreference(null)
() -> prepareSearch()
.setAllowPartialSearchResults(false)
.setPointInTime(new PointInTimeBuilder(pit))
.get()
Expand Down Expand Up @@ -365,7 +365,6 @@ public void testCanMatch() throws Exception {
assertResponse(
prepareSearch().setQuery(new RangeQueryBuilder("created_date").gte("2020-01-02").lte("2020-01-03"))
.setSearchType(SearchType.QUERY_THEN_FETCH)
.setPreference(null)
.setPreFilterShardSize(randomIntBetween(2, 3))
.setMaxConcurrentShardRequests(randomIntBetween(1, 2))
.setPointInTime(new PointInTimeBuilder(pitId)),
Expand Down Expand Up @@ -422,14 +421,14 @@ public void testPartialResults() throws Exception {
refresh();
String pitId = openPointInTime(new String[] { "test-*" }, TimeValue.timeValueMinutes(2));
try {
assertNoFailuresAndResponse(prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertNoFailuresAndResponse(prepareSearch().setPointInTime(new PointInTimeBuilder(pitId)), resp -> {
assertHitCount(resp, numDocs1 + numDocs2);
assertThat(resp.pointInTimeId(), equalTo(pitId));
});

internalCluster().restartNode(assignedNodeForIndex1);
assertResponse(
prepareSearch().setPreference(null).setAllowPartialSearchResults(true).setPointInTime(new PointInTimeBuilder(pitId)),
prepareSearch().setAllowPartialSearchResults(true).setPointInTime(new PointInTimeBuilder(pitId)),
resp -> {
assertFailures(resp);
assertThat(resp.pointInTimeId(), equalTo(pitId));
Expand Down Expand Up @@ -543,7 +542,7 @@ public void testOpenPITConcurrentShardRequests() throws Exception {
@SuppressWarnings({ "rawtypes", "unchecked" })
private void assertPagination(PointInTimeBuilder pit, int expectedNumDocs, int size, SortBuilder<?>... sorts) throws Exception {
Set<String> seen = new HashSet<>();
SearchRequestBuilder builder = prepareSearch().setPreference(null).setSize(size).setPointInTime(pit);
SearchRequestBuilder builder = prepareSearch().setSize(size).setPointInTime(pit);
for (SortBuilder<?> sort : sorts) {
builder.addSort(sort);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,6 @@ private void assertSearchSlicesWithPointInTime(String sliceField, String sortFie
int numSliceResults = 0;

SearchRequestBuilder request = prepareSearch().slice(new SliceBuilder(sliceField, id, numSlice))
.setPreference(null)
.setPointInTime(new PointInTimeBuilder(pointInTimeId))
.addSort(SortBuilders.fieldSort(sortField))
.setSize(randomIntBetween(10, 100));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,11 @@
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.client.internal.Client;
import org.elasticsearch.client.internal.ElasticsearchClient;
import org.elasticsearch.client.internal.FilterClient;
import org.elasticsearch.cluster.routing.Preference;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.search.builder.PointInTimeBuilder;

import java.util.Arrays;
import java.util.Random;
Expand Down Expand Up @@ -59,7 +61,8 @@ public RandomizingClient(Client client, Random random) {

@Override
public SearchRequestBuilder prepareSearch(String... indices) {
SearchRequestBuilder searchRequestBuilder = in.prepareSearch(indices)
SearchRequestBuilder searchRequestBuilder = new RandomizedSearchRequestBuilder(this)
.setIndices(indices)
.setSearchType(defaultSearchType)
.setPreference(defaultPreference)
.setBatchedReduceSize(batchedReduceSize);
Expand All @@ -84,4 +87,18 @@ public Client in() {
return super.in();
}

private class RandomizedSearchRequestBuilder extends SearchRequestBuilder {
RandomizedSearchRequestBuilder(ElasticsearchClient client) {
super(client);
}

@Override
public SearchRequestBuilder setPointInTime(PointInTimeBuilder pointInTimeBuilder) {
if (defaultPreference != null) {
setPreference(null);
}
return super.setPointInTime(pointInTimeBuilder);
}
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,7 @@ public void testRetryPointInTime() throws Exception {
final String pitId = client().execute(TransportOpenPointInTimeAction.TYPE, openPointInTimeRequest).actionGet().getPointInTimeId();
try {
assertNoFailuresAndResponse(
prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pitId)),
prepareSearch().setPointInTime(new PointInTimeBuilder(pitId)),
searchResponse -> {
assertThat(searchResponse.pointInTimeId(), equalTo(pitId));
assertHitCount(searchResponse, numDocs);
Expand All @@ -245,7 +245,6 @@ public void testRetryPointInTime() throws Exception {
assertNoFailuresAndResponse(
prepareSearch().setQuery(new RangeQueryBuilder("created_date").gte("2011-01-01").lte("2011-12-12"))
.setSearchType(SearchType.QUERY_THEN_FETCH)
.setPreference(null)
.setPreFilterShardSize(between(1, 10))
.setAllowPartialSearchResults(true)
.setPointInTime(new PointInTimeBuilder(pitId)),
Expand Down Expand Up @@ -286,7 +285,7 @@ public void testPointInTimeWithDeletedIndices() {
indicesAdmin().prepareDelete("index-1").get();
// Return partial results if allow partial search result is allowed
assertResponse(
prepareSearch().setPreference(null).setAllowPartialSearchResults(true).setPointInTime(new PointInTimeBuilder(pitId)),
prepareSearch().setAllowPartialSearchResults(true).setPointInTime(new PointInTimeBuilder(pitId)),
searchResponse -> {
assertFailures(searchResponse);
assertHitCount(searchResponse, index2);
Expand All @@ -295,7 +294,7 @@ public void testPointInTimeWithDeletedIndices() {
// Fails if allow partial search result is not allowed
expectThrows(
ElasticsearchException.class,
prepareSearch().setPreference(null).setAllowPartialSearchResults(false).setPointInTime(new PointInTimeBuilder(pitId))::get
prepareSearch().setAllowPartialSearchResults(false).setPointInTime(new PointInTimeBuilder(pitId))::get
);
} finally {
client().execute(TransportClosePointInTimeAction.TYPE, new ClosePointInTimeRequest(pitId)).actionGet();
Expand All @@ -321,7 +320,7 @@ public void testOpenPointInTimeWithNoIndexMatched() {
.getPointInTimeId();
try {
assertNoFailuresAndResponse(
prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pitId)),
prepareSearch().setPointInTime(new PointInTimeBuilder(pitId)),
searchResponse -> assertHitCount(searchResponse, numDocs)
);
} finally {
Expand All @@ -337,7 +336,7 @@ public void testOpenPointInTimeWithNoIndexMatched() {
.actionGet()
.getPointInTimeId();
try {
assertHitCountAndNoFailures(prepareSearch().setPreference(null).setPointInTime(new PointInTimeBuilder(pitId)), 0);
assertHitCountAndNoFailures(prepareSearch().setPointInTime(new PointInTimeBuilder(pitId)), 0);
} finally {
client().execute(TransportClosePointInTimeAction.TYPE, new ClosePointInTimeRequest(pitId)).actionGet();
}
Expand Down

0 comments on commit 93063b5

Please sign in to comment.