diff --git a/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java b/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java new file mode 100644 index 0000000000000..00013a8d396ba --- /dev/null +++ b/server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverRewriteIT.java @@ -0,0 +1,264 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0 and the Server Side Public License, v 1; you may not use this file except + * in compliance with, at your election, the Elastic License 2.0 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.search.retriever; + +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; +import org.elasticsearch.search.MockSearchService; +import org.elasticsearch.search.builder.SearchSourceBuilder; +import org.elasticsearch.test.ESIntegTestCase; +import org.elasticsearch.test.hamcrest.ElasticsearchAssertions; +import org.elasticsearch.xcontent.XContentBuilder; +import org.junit.Before; + +import java.io.IOException; +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() { + return List.of(MockSearchService.TestPlugin.class); + } + + private static String INDEX_DOCS = "docs"; + private static String INDEX_QUERIES = "queries"; + private static final String ID_FIELD = "_id"; + private static final String QUERY_FIELD = "query"; + + @Before + public void setup() throws Exception { + createIndex(INDEX_DOCS); + index(INDEX_DOCS, "doc_0", "{}"); + index(INDEX_DOCS, "doc_1", "{}"); + index(INDEX_DOCS, "doc_2", "{}"); + refresh(INDEX_DOCS); + + createIndex(INDEX_QUERIES); + index(INDEX_QUERIES, "query_0", "{ \"" + QUERY_FIELD + "\": \"doc_2\"}"); + index(INDEX_QUERIES, "query_1", "{ \"" + QUERY_FIELD + "\": \"doc_1\"}"); + index(INDEX_QUERIES, "query_2", "{ \"" + QUERY_FIELD + "\": \"doc_0\"}"); + refresh(INDEX_QUERIES); + } + + public void testRewrite() { + SearchSourceBuilder source = new SearchSourceBuilder(); + StandardRetrieverBuilder standard = new StandardRetrieverBuilder(); + standard.queryBuilder = QueryBuilders.termQuery(ID_FIELD, "doc_0"); + source.retriever(new AssertingRetrieverBuilder(standard)); + 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")); + }); + } + + public void testRewriteCompound() { + SearchSourceBuilder source = new SearchSourceBuilder(); + source.retriever(new AssertingCompoundRetrieverBuilder("query_0")); + 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("[open_point_in_time] action requires all shards to be available. Missing shards") + ); + } finally { + internalCluster().restartNode(randomDataNode); + } + } + + private static class AssertingRetrieverBuilder extends RetrieverBuilder { + private final RetrieverBuilder innerRetriever; + + private AssertingRetrieverBuilder(RetrieverBuilder innerRetriever) { + this.innerRetriever = innerRetriever; + } + + @Override + public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException { + assertNull(ctx.getPointInTimeBuilder()); + assertNull(ctx.convertToInnerHitsRewriteContext()); + assertNull(ctx.convertToCoordinatorRewriteContext()); + assertNull(ctx.convertToIndexMetadataContext()); + assertNull(ctx.convertToSearchExecutionContext()); + assertNull(ctx.convertToDataRewriteContext()); + var newRetriever = innerRetriever.rewrite(ctx); + if (newRetriever != innerRetriever) { + return new AssertingRetrieverBuilder(newRetriever); + } + return this; + } + + @Override + public void extractToSearchSourceBuilder(SearchSourceBuilder sourceBuilder, boolean compoundUsed) { + assertNull(sourceBuilder.retriever()); + innerRetriever.extractToSearchSourceBuilder(sourceBuilder, compoundUsed); + } + + @Override + public String getName() { + return "asserting"; + } + + @Override + protected void doToXContent(XContentBuilder builder, Params params) throws IOException {} + + @Override + protected boolean doEquals(Object o) { + return false; + } + + @Override + protected int doHashCode() { + return innerRetriever.doHashCode(); + } + } + + private static class AssertingCompoundRetrieverBuilder extends RetrieverBuilder { + private final String id; + private final SetOnce innerRetriever; + + private AssertingCompoundRetrieverBuilder(String id) { + this.id = id; + this.innerRetriever = new SetOnce<>(null); + } + + private AssertingCompoundRetrieverBuilder(String id, SetOnce innerRetriever) { + this.id = id; + this.innerRetriever = innerRetriever; + } + + @Override + public boolean isCompound() { + return true; + } + + @Override + public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException { + assertNotNull(ctx.getPointInTimeBuilder()); + assertNull(ctx.convertToInnerHitsRewriteContext()); + assertNull(ctx.convertToCoordinatorRewriteContext()); + assertNull(ctx.convertToIndexMetadataContext()); + assertNull(ctx.convertToSearchExecutionContext()); + assertNull(ctx.convertToDataRewriteContext()); + if (innerRetriever.get() != null) { + return this; + } + SetOnce innerRetriever = new SetOnce<>(); + ctx.registerAsyncAction((client, actionListener) -> { + SearchSourceBuilder source = new SearchSourceBuilder().pointInTimeBuilder(ctx.getPointInTimeBuilder()) + .query(QueryBuilders.termQuery(ID_FIELD, id)) + .fetchField(QUERY_FIELD); + client.search(new SearchRequest().source(source), new ActionListener<>() { + @Override + public void onResponse(SearchResponse response) { + String query = response.getHits().getAt(0).field(QUERY_FIELD).getValue(); + StandardRetrieverBuilder standard = new StandardRetrieverBuilder(); + standard.queryBuilder = QueryBuilders.termQuery(ID_FIELD, query); + innerRetriever.set(standard); + actionListener.onResponse(null); + } + + @Override + public void onFailure(Exception e) { + actionListener.onFailure(e); + } + }); + }); + return new AssertingCompoundRetrieverBuilder(id, innerRetriever); + } + + @Override + public void extractToSearchSourceBuilder(SearchSourceBuilder sourceBuilder, boolean compoundUsed) { + assertNull(sourceBuilder.retriever()); + innerRetriever.get().extractToSearchSourceBuilder(sourceBuilder, compoundUsed); + } + + @Override + public String getName() { + return "asserting"; + } + + @Override + protected void doToXContent(XContentBuilder builder, Params params) throws IOException { + throw new AssertionError("not implemented"); + } + + @Override + protected boolean doEquals(Object o) { + return false; + } + + @Override + protected int doHashCode() { + return id.hashCode(); + } + } +} diff --git a/server/src/main/java/org/elasticsearch/action/admin/indices/validate/query/TransportValidateQueryAction.java b/server/src/main/java/org/elasticsearch/action/admin/indices/validate/query/TransportValidateQueryAction.java index 1c3d168b8889d..577b34419e7c9 100644 --- a/server/src/main/java/org/elasticsearch/action/admin/indices/validate/query/TransportValidateQueryAction.java +++ b/server/src/main/java/org/elasticsearch/action/admin/indices/validate/query/TransportValidateQueryAction.java @@ -121,7 +121,11 @@ protected void doExecute(Task task, ValidateQueryRequest request, ActionListener if (request.query() == null) { rewriteListener.onResponse(request.query()); } else { - Rewriteable.rewriteAndFetch(request.query(), searchService.getRewriteContext(timeProvider, resolvedIndices), rewriteListener); + Rewriteable.rewriteAndFetch( + request.query(), + searchService.getRewriteContext(timeProvider, resolvedIndices, null), + rewriteListener + ); } } diff --git a/server/src/main/java/org/elasticsearch/action/explain/TransportExplainAction.java b/server/src/main/java/org/elasticsearch/action/explain/TransportExplainAction.java index bece9922f3e46..0379304a9901f 100644 --- a/server/src/main/java/org/elasticsearch/action/explain/TransportExplainAction.java +++ b/server/src/main/java/org/elasticsearch/action/explain/TransportExplainAction.java @@ -98,7 +98,7 @@ protected void doExecute(Task task, ExplainRequest request, ActionListener request.nowInMillis; - Rewriteable.rewriteAndFetch(request.query(), searchService.getRewriteContext(timeProvider, resolvedIndices), rewriteListener); + Rewriteable.rewriteAndFetch(request.query(), searchService.getRewriteContext(timeProvider, resolvedIndices, null), rewriteListener); } @Override diff --git a/server/src/main/java/org/elasticsearch/action/search/AbstractSearchAsyncAction.java b/server/src/main/java/org/elasticsearch/action/search/AbstractSearchAsyncAction.java index 3ff820c0e0354..4fd551994e2a0 100644 --- a/server/src/main/java/org/elasticsearch/action/search/AbstractSearchAsyncAction.java +++ b/server/src/main/java/org/elasticsearch/action/search/AbstractSearchAsyncAction.java @@ -709,7 +709,9 @@ public void sendSearchResponse(SearchResponseSections internalSearchResponse, At if (buildPointInTimeFromSearchResults()) { searchContextId = SearchContextId.encode(queryResults.asList(), aliasFilter, minTransportVersion); } else { - if (request.source() != null && request.source().pointInTimeBuilder() != null) { + if (request.source() != null + && request.source().pointInTimeBuilder() != null + && request.source().pointInTimeBuilder().singleSession() == false) { searchContextId = request.source().pointInTimeBuilder().getEncodedId(); } else { searchContextId = null; 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/TransportSearchAction.java b/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java index 84d233ec9710a..60a69dd6a7450 100644 --- a/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java +++ b/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java @@ -68,6 +68,7 @@ import org.elasticsearch.search.SearchPhaseResult; import org.elasticsearch.search.SearchService; import org.elasticsearch.search.aggregations.AggregationReduceContext; +import org.elasticsearch.search.builder.PointInTimeBuilder; import org.elasticsearch.search.builder.SearchSourceBuilder; import org.elasticsearch.search.internal.AliasFilter; import org.elasticsearch.search.internal.SearchContext; @@ -495,11 +496,72 @@ void executeRequest( } } }); + final SearchSourceBuilder source = original.source(); + if (shouldOpenPIT(source)) { + openPIT(client, original, searchService.getDefaultKeepAliveInMillis(), listener.delegateFailureAndWrap((delegate, resp) -> { + // We set the keep alive to -1 to indicate that we don't need the pit id in the response. + // This is needed since we delete the pit prior to sending the response so the id doesn't exist anymore. + source.pointInTimeBuilder(new PointInTimeBuilder(resp.getPointInTimeId()).setKeepAlive(TimeValue.MINUS_ONE)); + executeRequest(task, original, new ActionListener<>() { + @Override + public void onResponse(SearchResponse response) { + // we need to close the PIT first so we delay the release of the response to after the closing + response.incRef(); + closePIT( + client, + original.source().pointInTimeBuilder(), + () -> ActionListener.respondAndRelease(listener, response) + ); + } - Rewriteable.rewriteAndFetch( - original, - searchService.getRewriteContext(timeProvider::absoluteStartMillis, resolvedIndices), - rewriteListener + @Override + public void onFailure(Exception e) { + closePIT(client, original.source().pointInTimeBuilder(), () -> listener.onFailure(e)); + } + }, searchPhaseProvider); + })); + } else { + Rewriteable.rewriteAndFetch( + original, + searchService.getRewriteContext(timeProvider::absoluteStartMillis, resolvedIndices, original.pointInTimeBuilder()), + rewriteListener + ); + } + } + + /** + * Returns true if the provided source needs to open a shared point in time prior to executing the request. + */ + private boolean shouldOpenPIT(SearchSourceBuilder source) { + if (source == null) { + return false; + } + if (source.pointInTimeBuilder() != null) { + return false; + } + var retriever = source.retriever(); + return retriever != null && retriever.isCompound(); + } + + static void openPIT(Client client, SearchRequest request, long keepAliveMillis, ActionListener listener) { + OpenPointInTimeRequest pitReq = new OpenPointInTimeRequest(request.indices()).indicesOptions(request.indicesOptions()) + .preference(request.preference()) + .routing(request.routing()) + .keepAlive(TimeValue.timeValueMillis(keepAliveMillis)); + client.execute(TransportOpenPointInTimeAction.TYPE, pitReq, listener); + } + + static void closePIT(Client client, PointInTimeBuilder pit, Runnable next) { + client.execute( + TransportClosePointInTimeAction.TYPE, + new ClosePointInTimeRequest(pit.getEncodedId()), + ActionListener.runAfter(new ActionListener<>() { + @Override + public void onResponse(ClosePointInTimeResponse closePointInTimeResponse) {} + + @Override + public void onFailure(Exception e) {} + }, next) ); } diff --git a/server/src/main/java/org/elasticsearch/action/search/TransportSearchShardsAction.java b/server/src/main/java/org/elasticsearch/action/search/TransportSearchShardsAction.java index c3eea1fe557e7..21d9869b32747 100644 --- a/server/src/main/java/org/elasticsearch/action/search/TransportSearchShardsAction.java +++ b/server/src/main/java/org/elasticsearch/action/search/TransportSearchShardsAction.java @@ -115,7 +115,7 @@ protected void doExecute(Task task, SearchShardsRequest searchShardsRequest, Act Rewriteable.rewriteAndFetch( original, - searchService.getRewriteContext(timeProvider::absoluteStartMillis, resolvedIndices), + searchService.getRewriteContext(timeProvider::absoluteStartMillis, resolvedIndices, null), listener.delegateFailureAndWrap((delegate, searchRequest) -> { Index[] concreteIndices = resolvedIndices.getConcreteLocalIndices(); final Set indicesAndAliases = indexNameExpressionResolver.resolveExpressions(clusterState, searchRequest.indices()); diff --git a/server/src/main/java/org/elasticsearch/index/IndexService.java b/server/src/main/java/org/elasticsearch/index/IndexService.java index 413a31c0ccd6d..cbe94526819d9 100644 --- a/server/src/main/java/org/elasticsearch/index/IndexService.java +++ b/server/src/main/java/org/elasticsearch/index/IndexService.java @@ -799,6 +799,7 @@ public QueryRewriteContext newQueryRewriteContext( valuesSourceRegistry, allowExpensiveQueries, scriptService, + null, null ); } diff --git a/server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java b/server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java index 2a1062f8876d2..ac6512b0839e6 100644 --- a/server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java @@ -51,6 +51,7 @@ public CoordinatorRewriteContext( null, null, null, + null, null ); this.indexLongFieldRange = indexLongFieldRange; diff --git a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java index 81adfee36f923..9537aeec6a219 100644 --- a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java @@ -14,6 +14,7 @@ import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.regex.Regex; import org.elasticsearch.common.util.concurrent.CountDown; +import org.elasticsearch.core.Nullable; import org.elasticsearch.index.Index; import org.elasticsearch.index.IndexSettings; import org.elasticsearch.index.analysis.IndexAnalyzers; @@ -24,6 +25,7 @@ import org.elasticsearch.index.mapper.TextFieldMapper; import org.elasticsearch.script.ScriptCompiler; import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry; +import org.elasticsearch.search.builder.PointInTimeBuilder; import org.elasticsearch.xcontent.XContentParser; import org.elasticsearch.xcontent.XContentParserConfiguration; @@ -62,6 +64,7 @@ public class QueryRewriteContext { protected boolean mapUnmappedFieldAsString; protected Predicate allowedFields; private final ResolvedIndices resolvedIndices; + private final PointInTimeBuilder pit; public QueryRewriteContext( final XContentParserConfiguration parserConfiguration, @@ -77,7 +80,8 @@ public QueryRewriteContext( final ValuesSourceRegistry valuesSourceRegistry, final BooleanSupplier allowExpensiveQueries, final ScriptCompiler scriptService, - final ResolvedIndices resolvedIndices + final ResolvedIndices resolvedIndices, + final PointInTimeBuilder pit ) { this.parserConfiguration = parserConfiguration; @@ -95,6 +99,7 @@ public QueryRewriteContext( this.allowExpensiveQueries = allowExpensiveQueries; this.scriptService = scriptService; this.resolvedIndices = resolvedIndices; + this.pit = pit; } public QueryRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis) { @@ -112,6 +117,7 @@ public QueryRewriteContext(final XContentParserConfiguration parserConfiguration null, null, null, + null, null ); } @@ -120,7 +126,8 @@ public QueryRewriteContext( final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis, - final ResolvedIndices resolvedIndices + final ResolvedIndices resolvedIndices, + final PointInTimeBuilder pit ) { this( parserConfiguration, @@ -136,7 +143,8 @@ public QueryRewriteContext( null, null, null, - resolvedIndices + resolvedIndices, + pit ); } @@ -390,4 +398,12 @@ public Iterable> getAllFields() { public ResolvedIndices getResolvedIndices() { return resolvedIndices; } + + /** + * Returns the {@link PointInTimeBuilder} used by the search request, or null if not specified. + */ + @Nullable + public PointInTimeBuilder getPointInTimeBuilder() { + return pit; + } } diff --git a/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java b/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java index 4b16a5833cf41..f25a0c73ac25d 100644 --- a/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java @@ -269,6 +269,7 @@ private SearchExecutionContext( valuesSourceRegistry, allowExpensiveQueries, scriptService, + null, null ); this.shardId = shardId; diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesService.java b/server/src/main/java/org/elasticsearch/indices/IndicesService.java index 0d81d24e64646..3016530292766 100644 --- a/server/src/main/java/org/elasticsearch/indices/IndicesService.java +++ b/server/src/main/java/org/elasticsearch/indices/IndicesService.java @@ -137,6 +137,7 @@ import org.elasticsearch.repositories.RepositoriesService; import org.elasticsearch.script.ScriptService; import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry; +import org.elasticsearch.search.builder.PointInTimeBuilder; import org.elasticsearch.search.internal.AliasFilter; import org.elasticsearch.search.internal.SearchContext; import org.elasticsearch.search.internal.ShardSearchRequest; @@ -1755,8 +1756,8 @@ public AliasFilter buildAliasFilter(ClusterState state, String index, Set 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_partial_search_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..84a4eab897ba8 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,113 @@ 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); + searchRequest.scroll((Scroll) null); + ActionRequestValidationException validationErrors = searchRequest.validate(); + assertNotNull(validationErrors); + assertEquals(1, validationErrors.validationErrors().size()); + assertEquals( + "cannot specify a compound retriever and [allow_partial_search_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 "not_a_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); + searchRequest.scroll((Scroll) null); + 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; + } + })); + searchRequest.scroll((Scroll) null); + 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/action/search/TransportSearchActionTests.java b/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java index 98de321d792e0..edd253e945a9b 100644 --- a/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java @@ -1741,7 +1741,7 @@ protected void doWriteTo(StreamOutput out) throws IOException { NodeClient client = new NodeClient(settings, threadPool); SearchService searchService = mock(SearchService.class); - when(searchService.getRewriteContext(any(), any())).thenReturn(new QueryRewriteContext(null, null, null)); + when(searchService.getRewriteContext(any(), any(), any())).thenReturn(new QueryRewriteContext(null, null, null, null, null)); ClusterService clusterService = new ClusterService( settings, new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS), 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]")); } diff --git a/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java index 271df2a971fb1..2a3cc3a248f45 100644 --- a/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java +++ b/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java @@ -617,7 +617,8 @@ QueryRewriteContext createQueryRewriteContext() { null, () -> true, scriptService, - createMockResolvedIndices() + createMockResolvedIndices(), + null ); }