Add retriever to the query rewrite phase

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<Class<? extends Plugin>> 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<RetrieverBuilder> innerRetriever;
+
+        private AssertingCompoundRetrieverBuilder(String id) {
+            this.id = id;
+            this.innerRetriever = new SetOnce<>(null);
+        }
+
+        private AssertingCompoundRetrieverBuilder(String id, SetOnce<RetrieverBuilder> 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<RetrieverBuilder> 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();
+        }
+    }
+}

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
+            );
         }
     }
 

server/src/main/java/org/elasticsearch/action/explain/TransportExplainAction.java

@@ -98,7 +98,7 @@ protected void doExecute(Task task, ExplainRequest request, ActionListener<Expla
 
         assert request.query() != null;
         LongSupplier timeProvider = () -> request.nowInMillis;
-        Rewriteable.rewriteAndFetch(request.query(), searchService.getRewriteContext(timeProvider, resolvedIndices), rewriteListener);
+        Rewriteable.rewriteAndFetch(request.query(), searchService.getRewriteContext(timeProvider, resolvedIndices, null), rewriteListener);
     }
 
     @Override

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;

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/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<OpenPointInTimeResponse> 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)
         );
     }