Preserve final response headers in asynchronous search (#54349)

This change adds the response headers of the original search request
in the stored response in order to be able to restore them when retrieving a result
from the async-search index. It also ensures that response headers are preserved for
users that retrieve a final response on a running search task.
Partial response can eventually return response headers too but this change only ensures
that they are present when the response if final.

Relates #33936

x-pack/plugin/async-search/build.gradle

@@ -1,9 +1,3 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License;
- * you may not use this file except in compliance with the Elastic License.
- */
-
 evaluationDependsOn(xpackModule('core'))
 
 apply plugin: 'elasticsearch.esplugin'

x-pack/plugin/async-search/qa/rest/build.gradle

@@ -0,0 +1,27 @@
+import org.elasticsearch.gradle.info.BuildParams
+
+apply plugin: 'elasticsearch.testclusters'
+apply plugin: 'elasticsearch.esplugin'
+
+esplugin {
+  description 'Deprecated query plugin'
+  classname 'org.elasticsearch.query.DeprecatedQueryPlugin'
+}
+
+restResources {
+  restApi {
+    includeCore '_common', 'indices', 'index'
+    includeXpack 'async_search'
+  }
+}
+
+testClusters.integTest {
+  testDistribution = 'DEFAULT'
+  // add the deprecated query plugin
+  plugin file(project(':x-pack:plugin:async-search:qa:rest').tasks.bundlePlugin.archiveFile)
+  setting 'xpack.security.enabled', 'false'
+}
+
+test.enabled = false
+
+check.dependsOn integTest

x-pack/plugin/async-search/qa/rest/src/main/java/org/elasticsearch/query/DeprecatedQueryBuilder.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.query;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Query;
+import org.elasticsearch.common.ParsingException;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.logging.DeprecationLogger;
+import org.elasticsearch.common.xcontent.ObjectParser;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.index.query.AbstractQueryBuilder;
+import org.elasticsearch.index.query.QueryShardContext;
+
+import java.io.IOException;
+
+public class DeprecatedQueryBuilder extends AbstractQueryBuilder<DeprecatedQueryBuilder> {
+    private static final DeprecationLogger deprecationLogger = new DeprecationLogger(LogManager.getLogger("Deprecated"));
+
+    public static final String NAME = "deprecated";
+
+    public DeprecatedQueryBuilder() {}
+
+    DeprecatedQueryBuilder(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    protected void doWriteTo(StreamOutput out) {}
+
+    @Override
+    protected void doXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject(NAME);
+        builder.endObject();
+    }
+
+    private static final ObjectParser<DeprecatedQueryBuilder, Void> PARSER = new ObjectParser<>(NAME, DeprecatedQueryBuilder::new);
+
+    public static DeprecatedQueryBuilder fromXContent(XContentParser parser) {
+        try {
+            PARSER.apply(parser, null);
+            return new DeprecatedQueryBuilder();
+        } catch (IllegalArgumentException e) {
+            throw new ParsingException(parser.getTokenLocation(), e.getMessage(), e);
+        }
+    }
+
+    @Override
+    protected Query doToQuery(QueryShardContext context) {
+        deprecationLogger.deprecated("[deprecated] query");
+        return new MatchAllDocsQuery();
+    }
+
+    @Override
+    protected boolean doEquals(DeprecatedQueryBuilder other) {
+        return false;
+    }
+
+    @Override
+    protected int doHashCode() {
+        return 0;
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+}

x-pack/plugin/async-search/qa/rest/src/main/java/org/elasticsearch/query/DeprecatedQueryPlugin.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.query;
+
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.plugins.SearchPlugin;
+
+import java.util.List;
+
+import static java.util.Collections.singletonList;
+
+public class DeprecatedQueryPlugin extends Plugin implements SearchPlugin {
+
+    public DeprecatedQueryPlugin() {}
+
+    @Override
+    public List<QuerySpec<?>> getQueries() {
+        return singletonList(new QuerySpec<>("deprecated", DeprecatedQueryBuilder::new, p -> DeprecatedQueryBuilder.fromXContent(p)));
+    }
+}

x-pack/plugin/async-search/qa/rest/src/test/java/org/elasticsearch/qa/AsyncSearchRestIT.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.qa;
+
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
+import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase;
+
+public class AsyncSearchRestIT extends ESClientYamlSuiteTestCase {
+
+    public AsyncSearchRestIT(final ClientYamlTestCandidate testCandidate) {
+        super(testCandidate);
+    }
+
+    @ParametersFactory
+    public static Iterable<Object[]> parameters() throws Exception {
+        return ESClientYamlSuiteTestCase.createParameters();
+    }
+}

x-pack/plugin/async-search/qa/rest/src/test/resources/rest-api-spec/test/async-search/10_deprecation.yml

@@ -0,0 +1,104 @@
+setup:
+  - do:
+      indices.create:
+        index: test-1
+        body:
+          settings:
+            number_of_shards: "2"
+
+  - do:
+      indices.create:
+        index: test-2
+        body:
+          settings:
+            number_of_shards: "1"
+
+  - do:
+      indices.create:
+        index: test-3
+        body:
+          settings:
+            number_of_shards: "3"
+
+  - do:
+      index:
+        index:  test-2
+        body:   { max: 2 }
+
+  - do:
+      index:
+        index:  test-1
+        body:   { max: 1 }
+
+  - do:
+      index:
+        index:  test-3
+        body:   { max: 3 }
+
+  - do:
+      indices.refresh: {}
+
+---
+"Deprecation when retrieved from task":
+  - skip:
+      features: "warnings"
+
+  - do:
+      warnings:
+        - '[deprecated] query'
+      async_search.submit:
+        index: test-*
+        wait_for_completion_timeout: 10s
+        body:
+          query:
+            deprecated: {}
+
+  - is_false: id
+  - match:  { is_partial:                   false }
+  - length: { response.hits.hits:               3 }
+
+---
+"Deprecation when retrieved from store":
+  - skip:
+      features: "warnings"
+
+  - do:
+      warnings:
+        - '[deprecated] query'
+      async_search.submit:
+        index: test-*
+        wait_for_completion_timeout: 10s
+        keep_on_completion: true
+        body:
+          query:
+            deprecated: {}
+
+  - set:    { id:                              id }
+  - match:  { is_partial:                   false }
+  - length: { response.hits.hits:               3 }
+
+  - do:
+      warnings:
+        - '[deprecated] query'
+      async_search.get:
+        id: "$id"
+
+  - match:  { is_partial:                       false }
+  - length: { response.hits.hits:                   3 }
+
+  - do:
+      async_search.delete:
+        id: "$id"
+
+  - match: { acknowledged:   true }
+
+  - do:
+      catch: missing
+      async_search.get:
+        id: "$id"
+
+  - do:
+      catch: missing
+      async_search.delete:
+        id: "$id"
+

x-pack/plugin/async-search/src/main/java/org/elasticsearch/xpack/search/AsyncSearchIndexService.java

@@ -47,6 +47,7 @@
 import java.util.Base64;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
@@ -63,6 +64,7 @@ class AsyncSearchIndexService {
     public static final String INDEX = ".async-search";
 
     public static final String HEADERS_FIELD = "headers";
+    public static final String RESPONSE_HEADERS_FIELD = "response_headers";
     public static final String EXPIRATION_TIME_FIELD = "expiration_time";
     public static final String RESULT_FIELD = "result";
 
@@ -86,6 +88,10 @@ private static XContentBuilder mappings() throws IOException {
                             .field("type", "object")
                             .field("enabled", "false")
                         .endObject()
+                        .startObject(RESPONSE_HEADERS_FIELD)
+                            .field("type", "object")
+                            .field("enabled", "false")
+                        .endObject()
                         .startObject(RESULT_FIELD)
                             .field("type", "object")
                             .field("enabled", "false")
@@ -172,9 +178,11 @@ void storeInitialResponse(String docId,
      * Stores the final response if the place-holder document is still present (update).
      */
     void storeFinalResponse(String docId,
+                            Map<String, List<String>> responseHeaders,
                             AsyncSearchResponse response,
                             ActionListener<UpdateResponse> listener) throws IOException {
         Map<String, Object> source = new HashMap<>();
+        source.put(RESPONSE_HEADERS_FIELD, responseHeaders);
         source.put(RESULT_FIELD, encodeResponse(response));
         UpdateRequest request = new UpdateRequest()
             .index(INDEX)
@@ -249,8 +257,11 @@ AsyncSearchTask getTask(TaskManager taskManager, AsyncSearchId searchId) throws
     /**
      * Gets the response from the index if present, or delegate a {@link ResourceNotFoundException}
      * failure to the provided listener if not.
+     * When the provided <code>restoreResponseHeaders</code> is <code>true</code>, this method also restores the
+     * response headers of the original request in the current thread context.
      */
     void getResponse(AsyncSearchId searchId,
+                     boolean restoreResponseHeaders,
                      ActionListener<AsyncSearchResponse> listener) {
         final Authentication current = securityContext.getAuthentication();
         GetRequest internalGet = new GetRequest(INDEX)
@@ -271,6 +282,12 @@ void getResponse(AsyncSearchId searchId,
                     return;
                 }
 
+                if (restoreResponseHeaders) {
+                    @SuppressWarnings("unchecked")
+                    Map<String, List<String>> responseHeaders = (Map<String, List<String>>) get.getSource().get(RESPONSE_HEADERS_FIELD);
+                    restoreResponseHeadersContext(securityContext.getThreadContext(), responseHeaders);
+                }
+
                 String encoded = (String) get.getSource().get(RESULT_FIELD);
                 listener.onResponse(encoded != null ? decodeResponse(encoded) : null);
             },
@@ -339,4 +356,15 @@ AsyncSearchResponse decodeResponse(String value) throws IOException {
             }
         }
     }
+
+    /**
+     * Restores the provided <code>responseHeaders</code> to the current thread context.
+     */
+    static void restoreResponseHeadersContext(ThreadContext threadContext, Map<String, List<String>> responseHeaders) {
+        for (Map.Entry<String, List<String>> entry : responseHeaders.entrySet()) {
+            for (String value : entry.getValue()) {
+                threadContext.addResponseHeader(entry.getKey(), value);
+            }
+        }
+    }
 }