From 29266a7b468e5834e4d08fb3c8a7781b195f9639 Mon Sep 17 00:00:00 2001 From: Panagiotis Bailis Date: Thu, 25 Jul 2024 09:41:52 +0300 Subject: [PATCH] ensuring that allow_search_partial_results is disabled for compound retrievers --- .../search/retriever/RetrieverRewriteIT.java | 49 ++++++++ .../action/search/SearchPhase.java | 2 +- .../action/search/SearchRequest.java | 3 +- .../TransportOpenPointInTimeAction.java | 24 ++++ .../search/builder/SearchSourceBuilder.java | 14 ++- .../action/search/SearchRequestTests.java | 106 ++++++++++++++++++ .../retriever/RetrieverBuilderErrorTests.java | 16 +-- 7 files changed, 202 insertions(+), 12 deletions(-) diff --git a/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java b/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java index 8f28e5f279cf6..cfe3a13bb26c5 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java +++ b/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java @@ -11,9 +11,16 @@ import org.apache.lucene.search.TotalHits; import org.apache.lucene.util.SetOnce; import org.elasticsearch.action.ActionListener; +import org.elasticsearch.action.admin.cluster.health.ClusterHealthResponse; +import org.elasticsearch.action.search.SearchPhaseExecutionException; import org.elasticsearch.action.search.SearchRequest; import org.elasticsearch.action.search.SearchRequestBuilder; import org.elasticsearch.action.search.SearchResponse; +import org.elasticsearch.cluster.health.ClusterHealthStatus; +import org.elasticsearch.cluster.node.DiscoveryNode; +import org.elasticsearch.cluster.node.DiscoveryNodeRole; +import org.elasticsearch.common.Priority; +import org.elasticsearch.common.settings.Settings; import org.elasticsearch.index.query.QueryBuilders; import org.elasticsearch.index.query.QueryRewriteContext; import org.elasticsearch.plugins.Plugin; @@ -28,8 +35,12 @@ import java.util.Collection; import java.util.List; +import static org.elasticsearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_REPLICAS; +import static org.elasticsearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_SHARDS; +import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; +@ESIntegTestCase.ClusterScope(numDataNodes = 3) public class RetrieverRewriteIT extends ESIntegTestCase { @Override protected Collection> nodePlugins() { @@ -64,6 +75,7 @@ public void testRewrite() { SearchRequestBuilder req = client().prepareSearch(INDEX_DOCS, INDEX_QUERIES).setSource(source); ElasticsearchAssertions.assertResponse(req, resp -> { assertNull(resp.pointInTimeId()); + assertNotNull(resp.getHits().getTotalHits()); assertThat(resp.getHits().getTotalHits().value, equalTo(1L)); assertThat(resp.getHits().getTotalHits().relation, equalTo(TotalHits.Relation.EQUAL_TO)); assertThat(resp.getHits().getAt(0).getId(), equalTo("doc_0")); @@ -76,12 +88,49 @@ public void testRewriteCompound() { SearchRequestBuilder req = client().prepareSearch(INDEX_DOCS, INDEX_QUERIES).setSource(source); ElasticsearchAssertions.assertResponse(req, resp -> { assertNull(resp.pointInTimeId()); + assertNotNull(resp.getHits().getTotalHits()); assertThat(resp.getHits().getTotalHits().value, equalTo(1L)); assertThat(resp.getHits().getTotalHits().relation, equalTo(TotalHits.Relation.EQUAL_TO)); assertThat(resp.getHits().getAt(0).getId(), equalTo("doc_2")); }); } + public void testRewriteCompoundRetrieverShouldThrowForPartialResults() throws Exception { + final String testIndex = "test"; + createIndex(testIndex, Settings.builder().put(SETTING_NUMBER_OF_SHARDS, 3).put(SETTING_NUMBER_OF_REPLICAS, 0).build()); + for (int i = 0; i < 50; i++) { + index(testIndex, "doc_" + i, "{}"); + } + refresh(testIndex); + + SearchSourceBuilder source = new SearchSourceBuilder(); + source.retriever(new AssertingCompoundRetrieverBuilder("doc_0")); + final String randomDataNode = internalCluster().getNodeNameThat( + settings -> DiscoveryNode.hasRole(settings, DiscoveryNodeRole.DATA_ROLE) + ); + try { + ensureGreen(testIndex); + if (false == internalCluster().stopNode(randomDataNode)) { + throw new IllegalStateException("node did not stop"); + } + assertBusy(() -> { + ClusterHealthResponse healthResponse = clusterAdmin().prepareHealth(testIndex) + .setWaitForStatus(ClusterHealthStatus.RED) // we are now known red because the primary shard is missing + .setWaitForEvents(Priority.LANGUID) // ensures that the update has occurred + .execute() + .actionGet(); + assertThat(healthResponse.getStatus(), equalTo(ClusterHealthStatus.RED)); + }); + SearchPhaseExecutionException ex = expectThrows( + SearchPhaseExecutionException.class, + client().prepareSearch(testIndex).setSource(source)::get + ); + assertThat(ex.getDetailedMessage(), containsString("Cannot execute [open_point_in_time] action due to missing shards")); + } finally { + internalCluster().restartNode(randomDataNode); + } + } + private static class AssertingRetrieverBuilder extends RetrieverBuilder { private final RetrieverBuilder innerRetriever; diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchPhase.java b/server/src/main/java/org/elasticsearch/action/search/SearchPhase.java index 7ad81154691c0..4e46858123c03 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchPhase.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchPhase.java @@ -42,7 +42,7 @@ public void start() { } } - static void doCheckNoMissingShards(String phaseName, SearchRequest request, GroupShardsIterator shardsIts) { + protected void doCheckNoMissingShards(String phaseName, SearchRequest request, GroupShardsIterator shardsIts) { assert request.allowPartialSearchResults() != null : "SearchRequest missing setting for allowPartialSearchResults"; if (request.allowPartialSearchResults() == false) { final StringBuilder missingShards = new StringBuilder(); diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchRequest.java b/server/src/main/java/org/elasticsearch/action/search/SearchRequest.java index 514e8d10eeca1..415a2ba7da9f8 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchRequest.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchRequest.java @@ -320,9 +320,10 @@ public void writeTo(StreamOutput out) throws IOException { public ActionRequestValidationException validate() { ActionRequestValidationException validationException = null; boolean scroll = scroll() != null; + boolean allowPartialSearchResults = allowPartialSearchResults() != null && allowPartialSearchResults(); if (source != null) { - validationException = source.validate(validationException, scroll); + validationException = source.validate(validationException, scroll, allowPartialSearchResults); } if (scroll) { if (requestCache != null && requestCache) { diff --git a/server/src/main/java/org/elasticsearch/action/search/TransportOpenPointInTimeAction.java b/server/src/main/java/org/elasticsearch/action/search/TransportOpenPointInTimeAction.java index a31593d06a521..0019ab0573c6e 100644 --- a/server/src/main/java/org/elasticsearch/action/search/TransportOpenPointInTimeAction.java +++ b/server/src/main/java/org/elasticsearch/action/search/TransportOpenPointInTimeAction.java @@ -217,6 +217,30 @@ SearchPhase openPointInTimePhase( searchRequest.getMaxConcurrentShardRequests(), clusters ) { + @Override + protected void doCheckNoMissingShards( + String phaseName, + SearchRequest request, + GroupShardsIterator shardsIts + ) { + final StringBuilder missingShards = new StringBuilder(); + // Fail-fast verification of all shards being available + for (int index = 0; index < shardsIts.size(); index++) { + final SearchShardIterator shardRoutings = shardsIts.get(index); + if (shardRoutings.size() == 0) { + if (missingShards.isEmpty() == false) { + missingShards.append(", "); + } + missingShards.append(shardRoutings.shardId()); + } + } + if (missingShards.isEmpty() == false) { + // Status red - shard is missing all copies and would produce partial results for an index search + final String msg = "Cannot execute [open_point_in_time] action due to missing shards [" + missingShards + "]."; + throw new SearchPhaseExecutionException(phaseName, msg, null, ShardSearchFailure.EMPTY_ARRAY); + } + } + @Override protected void executePhaseOnShard( SearchShardIterator shardIt, diff --git a/server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java b/server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java index a80378db99348..2cab29a7bc289 100644 --- a/server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java @@ -2194,14 +2194,24 @@ public boolean supportsParallelCollection(ToLongFunction fieldCardinalit } private void validate() throws ValidationException { - var exceptions = validate(null, false); + var exceptions = validate(null, false, false); if (exceptions != null) { throw exceptions; } } - public ActionRequestValidationException validate(ActionRequestValidationException validationException, boolean isScroll) { + public ActionRequestValidationException validate( + ActionRequestValidationException validationException, + boolean isScroll, + boolean allowPartialSearchResults + ) { if (retriever() != null) { + if (allowPartialSearchResults && retriever().isCompound()) { + validationException = addValidationError( + "cannot specify a compound retriever and [allow_search_partial_results]", + validationException + ); + } List specified = new ArrayList<>(); if (subSearches().isEmpty() == false) { specified.add(QUERY_FIELD.getPreferredName()); diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchRequestTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchRequestTests.java index 475f44238f36e..63dd47d8d51d2 100644 --- a/server/src/test/java/org/elasticsearch/action/search/SearchRequestTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchRequestTests.java @@ -29,6 +29,7 @@ import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder; import org.elasticsearch.search.rank.TestRankBuilder; import org.elasticsearch.search.rescore.QueryRescorerBuilder; +import org.elasticsearch.search.retriever.RetrieverBuilder; import org.elasticsearch.search.slice.SliceBuilder; import org.elasticsearch.search.suggest.SuggestBuilder; import org.elasticsearch.search.suggest.term.TermSuggestionBuilder; @@ -37,6 +38,7 @@ import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.TransportVersionUtils; import org.elasticsearch.test.VersionUtils; +import org.elasticsearch.xcontent.XContentBuilder; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -257,6 +259,110 @@ public void testValidate() throws IOException { assertEquals(1, validationErrors.validationErrors().size()); assertEquals("cannot use `collapse` in a scroll context", validationErrors.validationErrors().get(0)); } + { + // allow_partial_results and compound retriever + SearchRequest searchRequest = createSearchRequest().source(new SearchSourceBuilder().retriever(new RetrieverBuilder() { + @Override + public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder, boolean compoundUsed) { + // no-op + } + + @Override + public String getName() { + return "compound_retriever"; + } + + @Override + protected void doToXContent(XContentBuilder builder, Params params) throws IOException {} + + @Override + protected boolean doEquals(Object o) { + return false; + } + + @Override + protected int doHashCode() { + return 0; + } + + @Override + public boolean isCompound() { + return true; + } + })); + searchRequest.allowPartialSearchResults(true); + ActionRequestValidationException validationErrors = searchRequest.validate(); + assertNotNull(validationErrors); + assertEquals(1, validationErrors.validationErrors().size()); + assertEquals( + "cannot specify a compound retriever and [allow_search_partial_results]", + validationErrors.validationErrors().get(0) + ); + } + { + // allow_partial_results and non-compound retriever + SearchRequest searchRequest = createSearchRequest().source(new SearchSourceBuilder().retriever(new RetrieverBuilder() { + @Override + public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder, boolean compoundUsed) { + // no-op + } + + @Override + public String getName() { + return "compound_retriever"; + } + + @Override + protected void doToXContent(XContentBuilder builder, Params params) throws IOException {} + + @Override + protected boolean doEquals(Object o) { + return false; + } + + @Override + protected int doHashCode() { + return 0; + } + })); + searchRequest.allowPartialSearchResults(true); + ActionRequestValidationException validationErrors = searchRequest.validate(); + assertNull(validationErrors); + } + { + // allow_partial_results not defined and compound retriever + SearchRequest searchRequest = new SearchRequest().source(new SearchSourceBuilder().retriever(new RetrieverBuilder() { + @Override + public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder, boolean compoundUsed) { + // no-op + } + + @Override + public String getName() { + return "compound_retriever"; + } + + @Override + protected void doToXContent(XContentBuilder builder, Params params) throws IOException {} + + @Override + protected boolean doEquals(Object o) { + return false; + } + + @Override + protected int doHashCode() { + return 0; + } + + @Override + public boolean isCompound() { + return true; + } + })); + ActionRequestValidationException validationErrors = searchRequest.validate(); + assertNull(validationErrors); + } { // search_after and `from` isn't valid SearchRequest searchRequest = createSearchRequest().source(new SearchSourceBuilder()); diff --git a/server/src/test/java/org/elasticsearch/search/retriever/RetrieverBuilderErrorTests.java b/server/src/test/java/org/elasticsearch/search/retriever/RetrieverBuilderErrorTests.java index 97f7c018974fa..b4eb1e31176bc 100644 --- a/server/src/test/java/org/elasticsearch/search/retriever/RetrieverBuilderErrorTests.java +++ b/server/src/test/java/org/elasticsearch/search/retriever/RetrieverBuilderErrorTests.java @@ -36,7 +36,7 @@ public void testRetrieverExtractionErrors() throws IOException { ) { SearchSourceBuilder ssb = new SearchSourceBuilder(); ssb.parseXContent(parser, true, nf -> true); - ActionRequestValidationException iae = ssb.validate(null, false); + ActionRequestValidationException iae = ssb.validate(null, false, false); assertNotNull(iae); assertThat(iae.getMessage(), containsString("cannot specify [retriever] and [query]")); } @@ -50,7 +50,7 @@ public void testRetrieverExtractionErrors() throws IOException { ) { SearchSourceBuilder ssb = new SearchSourceBuilder(); ssb.parseXContent(parser, true, nf -> true); - ActionRequestValidationException iae = ssb.validate(null, false); + ActionRequestValidationException iae = ssb.validate(null, false, false); assertNotNull(iae); assertThat(iae.getMessage(), containsString("cannot specify [retriever] and [knn]")); } @@ -58,7 +58,7 @@ public void testRetrieverExtractionErrors() throws IOException { try (XContentParser parser = createParser(JsonXContent.jsonXContent, "{\"search_after\": [1], \"retriever\":{\"standard\":{}}}")) { SearchSourceBuilder ssb = new SearchSourceBuilder(); ssb.parseXContent(parser, true, nf -> true); - ActionRequestValidationException iae = ssb.validate(null, false); + ActionRequestValidationException iae = ssb.validate(null, false, false); assertNotNull(iae); assertThat(iae.getMessage(), containsString("cannot specify [retriever] and [search_after]")); @@ -67,7 +67,7 @@ public void testRetrieverExtractionErrors() throws IOException { try (XContentParser parser = createParser(JsonXContent.jsonXContent, "{\"terminate_after\": 1, \"retriever\":{\"standard\":{}}}")) { SearchSourceBuilder ssb = new SearchSourceBuilder(); ssb.parseXContent(parser, true, nf -> true); - ActionRequestValidationException iae = ssb.validate(null, false); + ActionRequestValidationException iae = ssb.validate(null, false, false); assertNotNull(iae); assertThat(iae.getMessage(), containsString("cannot specify [retriever] and [terminate_after]")); } @@ -75,7 +75,7 @@ public void testRetrieverExtractionErrors() throws IOException { try (XContentParser parser = createParser(JsonXContent.jsonXContent, "{\"sort\": [\"field\"], \"retriever\":{\"standard\":{}}}")) { SearchSourceBuilder ssb = new SearchSourceBuilder(); ssb.parseXContent(parser, true, nf -> true); - ActionRequestValidationException iae = ssb.validate(null, false); + ActionRequestValidationException iae = ssb.validate(null, false, false); assertNotNull(iae); assertThat(iae.getMessage(), containsString("cannot specify [retriever] and [sort]")); } @@ -88,7 +88,7 @@ public void testRetrieverExtractionErrors() throws IOException { ) { SearchSourceBuilder ssb = new SearchSourceBuilder(); ssb.parseXContent(parser, true, nf -> true); - ActionRequestValidationException iae = ssb.validate(null, false); + ActionRequestValidationException iae = ssb.validate(null, false, false); assertNotNull(iae); assertThat(iae.getMessage(), containsString("cannot specify [retriever] and [rescore]")); } @@ -96,7 +96,7 @@ public void testRetrieverExtractionErrors() throws IOException { try (XContentParser parser = createParser(JsonXContent.jsonXContent, "{\"min_score\": 2, \"retriever\":{\"standard\":{}}}")) { SearchSourceBuilder ssb = new SearchSourceBuilder(); ssb.parseXContent(parser, true, nf -> true); - ActionRequestValidationException iae = ssb.validate(null, false); + ActionRequestValidationException iae = ssb.validate(null, false, false); assertNotNull(iae); assertThat(iae.getMessage(), containsString("cannot specify [retriever] and [min_score]")); } @@ -109,7 +109,7 @@ public void testRetrieverExtractionErrors() throws IOException { ) { SearchSourceBuilder ssb = new SearchSourceBuilder(); ssb.parseXContent(parser, true, nf -> true); - ActionRequestValidationException iae = ssb.validate(null, false); + ActionRequestValidationException iae = ssb.validate(null, false, false); assertNotNull(iae); assertThat(iae.getMessage(), containsString("cannot specify [retriever] and [query, terminate_after, min_score]")); }