ensuring that allow_search_partial_results is disabled for compound r…

…etrievers

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<Class<? extends Plugin>> 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;
 

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<SearchShardIterator> shardsIts) {
+    protected void doCheckNoMissingShards(String phaseName, SearchRequest request, GroupShardsIterator<SearchShardIterator> shardsIts) {
         assert request.allowPartialSearchResults() != null : "SearchRequest missing setting for allowPartialSearchResults";
         if (request.allowPartialSearchResults() == false) {
             final StringBuilder missingShards = new StringBuilder();

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) {

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<SearchShardIterator> 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,

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -2194,14 +2194,24 @@ public boolean supportsParallelCollection(ToLongFunction<String> 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<String> specified = new ArrayList<>();
             if (subSearches().isEmpty() == false) {
                 specified.add(QUERY_FIELD.getPreferredName());

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());

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,15 +50,15 @@ 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]"));
         }
 
         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,15 +67,15 @@ 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]"));
         }
 
         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,15 +88,15 @@ 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]"));
         }
 
         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]"));
         }