Fix multi-search with template doesn't return status code (opensearch…

…-project#16265) * Fix multi-search with template doesn't return status code Signed-off-by: Gao Binlong <[email protected]> * Modify changelog Signed-off-by: Gao Binlong <[email protected]> --------- Signed-off-by: Gao Binlong <[email protected]>
dk2k · Oct 21, 2024 · 4b9a024 · 4b9a024
1 parent 35b2023
commit 4b9a024
Show file tree

Hide file tree

Showing 6 changed files with 136 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -73,6 +73,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Streaming bulk request hangs ([#16158](https://github.com/opensearch-project/OpenSearch/pull/16158))
 - Fix warnings from SLF4J on startup when repository-s3 is installed ([#16194](https://github.com/opensearch-project/OpenSearch/pull/16194))
 - Fix protobuf-java leak through client library dependencies ([#16254](https://github.com/opensearch-project/OpenSearch/pull/16254))
+- Fix multi-search with template doesn't return status code ([#16265](https://github.com/opensearch-project/OpenSearch/pull/16265))
 
 ### Security
 

diff --git a/...ng-mustache/src/main/java/org/opensearch/script/mustache/MultiSearchTemplateResponse.java b/...ng-mustache/src/main/java/org/opensearch/script/mustache/MultiSearchTemplateResponse.java
@@ -32,6 +32,7 @@
 
 package org.opensearch.script.mustache;
 
+import org.opensearch.ExceptionsHelper;
 import org.opensearch.OpenSearchException;
 import org.opensearch.action.search.MultiSearchResponse;
 import org.opensearch.common.Nullable;
@@ -167,6 +168,7 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
             if (item.isFailure()) {
                 builder.startObject();
                 OpenSearchException.generateFailureXContent(builder, params, item.getFailure(), true);
+                builder.field(Fields.STATUS, ExceptionsHelper.status(item.getFailure()).getStatus());
                 builder.endObject();
             } else {
                 item.getResponse().toXContent(builder, params);
@@ -179,6 +181,7 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
 
     static final class Fields {
         static final String RESPONSES = "responses";
+        static final String STATUS = "status";
     }
 
     public static MultiSearchTemplateResponse fromXContext(XContentParser parser) {

diff --git a/...es/lang-mustache/src/main/java/org/opensearch/script/mustache/SearchTemplateResponse.java b/...es/lang-mustache/src/main/java/org/opensearch/script/mustache/SearchTemplateResponse.java
@@ -120,7 +120,10 @@ public static SearchTemplateResponse fromXContent(XContentParser parser) throws
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         if (hasResponse()) {
-            response.toXContent(builder, params);
+            builder.startObject();
+            response.innerToXContent(builder, params);
+            builder.field(MultiSearchTemplateResponse.Fields.STATUS, response.status().getStatus());
+            builder.endObject();
         } else {
             builder.startObject();
             // we can assume the template is always json as we convert it before compiling it

diff --git a/...stache/src/test/java/org/opensearch/script/mustache/MultiSearchTemplateResponseTests.java b/...stache/src/test/java/org/opensearch/script/mustache/MultiSearchTemplateResponseTests.java
@@ -34,8 +34,12 @@
 import org.opensearch.OpenSearchException;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.action.search.ShardSearchFailure;
+import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.core.common.Strings;
+import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.core.xcontent.ToXContent;
+import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.search.internal.InternalSearchResponse;
 import org.opensearch.test.AbstractXContentTestCase;
@@ -44,6 +48,7 @@
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
+import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertToXContentEquivalent;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.nullValue;
@@ -177,4 +182,78 @@ public void testFromXContentWithFailures() throws IOException {
             ToXContent.EMPTY_PARAMS
         );
     }
+
+    public void testToXContentWithFailures() throws Exception {
+        long overallTookInMillis = randomNonNegativeLong();
+        MultiSearchTemplateResponse.Item[] items = new MultiSearchTemplateResponse.Item[2];
+
+        long tookInMillis = 1L;
+        int totalShards = 2;
+        int successfulShards = 2;
+        int skippedShards = 0;
+        InternalSearchResponse internalSearchResponse = InternalSearchResponse.empty();
+        SearchTemplateResponse searchTemplateResponse = new SearchTemplateResponse();
+        SearchResponse searchResponse = new SearchResponse(
+            internalSearchResponse,
+            null,
+            totalShards,
+            successfulShards,
+            skippedShards,
+            tookInMillis,
+            ShardSearchFailure.EMPTY_ARRAY,
+            SearchResponse.Clusters.EMPTY
+        );
+        searchTemplateResponse.setResponse(searchResponse);
+        items[0] = new MultiSearchTemplateResponse.Item(searchTemplateResponse, null);
+
+        items[1] = new MultiSearchTemplateResponse.Item(null, new IllegalArgumentException("Invalid argument"));
+
+        MultiSearchTemplateResponse response = new MultiSearchTemplateResponse(items, overallTookInMillis);
+
+        XContentType contentType = randomFrom(XContentType.values());
+        XContentBuilder expectedResponse = MediaTypeRegistry.contentBuilder(contentType)
+            .startObject()
+            .field("took", overallTookInMillis)
+            .startArray("responses")
+            .startObject()
+            .field("took", 1)
+            .field("timed_out", false)
+            .startObject("_shards")
+            .field("total", 2)
+            .field("successful", 2)
+            .field("skipped", 0)
+            .field("failed", 0)
+            .endObject()
+            .startObject("hits")
+            .startObject("total")
+            .field("value", 0)
+            .field("relation", "eq")
+            .endObject()
+            .field("max_score", 0.0F)
+            .startArray("hits")
+            .endArray()
+            .endObject()
+            .field("status", 200)
+            .endObject()
+            .startObject()
+            .startObject("error")
+            .field("type", "illegal_argument_exception")
+            .field("reason", "Invalid argument")
+            .startArray("root_cause")
+            .startObject()
+            .field("type", "illegal_argument_exception")
+            .field("reason", "Invalid argument")
+            .endObject()
+            .endArray()
+            .endObject()
+            .field("status", 400)
+            .endObject()
+            .endArray()
+            .endObject();
+
+        XContentBuilder actualResponse = MediaTypeRegistry.contentBuilder(contentType);
+        response.toXContent(actualResponse, ToXContent.EMPTY_PARAMS);
+
+        assertToXContentEquivalent(BytesReference.bytes(expectedResponse), BytesReference.bytes(actualResponse), contentType);
+    }
 }
diff --git a/...ng-mustache/src/test/java/org/opensearch/script/mustache/SearchTemplateResponseTests.java b/...ng-mustache/src/test/java/org/opensearch/script/mustache/SearchTemplateResponseTests.java
@@ -234,6 +234,7 @@ public void testSearchResponseToXContent() throws IOException {
             .endObject()
             .endArray()
             .endObject()
+            .field("status", 200)
             .endObject();
 
         XContentBuilder actualResponse = MediaTypeRegistry.contentBuilder(contentType);

diff --git a/.../src/yamlRestTest/resources/rest-api-spec/test/lang_mustache/50_multi_search_template.yml b/.../src/yamlRestTest/resources/rest-api-spec/test/lang_mustache/50_multi_search_template.yml
@@ -205,3 +205,51 @@ setup:
   - match:  { responses.1.hits.total.relation:  eq }
   - match:  { responses.2.hits.total.value:     1  }
   - match:  { responses.1.hits.total.relation:  eq }
+
+---
+"Test multi-search template returns status code":
+  - skip:
+      version: " - 2.17.99"
+      reason: Fixed in 2.18.0.
+  - do:
+      msearch_template:
+        rest_total_hits_as_int: true
+        body:
+          # Search 0 is OK
+          - index: index_*
+          - source: '{"query": {"match": {"foo": "{{value}}"} } }'
+            params:
+              value: "foo"
+          # Search 1 has an unclosed JSON template
+          - index: index_*
+          - source: '{"query": {"match": {'
+            params:
+              field: "foo"
+              value: "bar"
+          # Search 2 is OK
+          - index: _all
+          - source:
+              query:
+                match: {foo: "{{text}}"}
+            params:
+              text: "baz"
+          # Search 3 has an unknown query type
+          - index: index_*
+          - source: '{"query": {"{{query_type}}": {} }' # Unknown query type
+            params:
+              query_type: "unknown"
+          # Search 4 has an unsupported track_total_hits
+          - index: index_*
+          - source: '{"query": {"match_all": {} }, "track_total_hits": "{{trackTotalHits}}" }'
+            params:
+              trackTotalHits: 1
+          # Search 5 has an unknown index
+          - index: unknown_index
+          - source: '{"query": {"match_all": {} } }'
+
+  - match:  { responses.0.status:     200  }
+  - match:  { responses.1.status:     400  }
+  - match:  { responses.2.status:     200  }
+  - match:  { responses.3.status:     400  }
+  - match:  { responses.4.status:     400  }
+  - match:  { responses.5.status:     404  }