Add SearchExtBuilders to SearchResponse. [Issue #9328](#9328)

Signed-off-by: Austin Lee <austin@aryn.ai>
opensearch-project · Aug 16, 2023 · 15857c3 · 15857c3
1 parent b58bec4
commit 15857c3
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 1 deletion.
diff --git a/server/src/main/java/org/opensearch/action/search/SearchResponse.java b/server/src/main/java/org/opensearch/action/search/SearchResponse.java
@@ -49,6 +49,7 @@
 import org.opensearch.core.xcontent.XContentParser.Token;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.rest.action.RestActions;
+import org.opensearch.search.SearchExtBuilder;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
 import org.opensearch.search.aggregations.Aggregations;
@@ -80,6 +81,7 @@ public class SearchResponse extends ActionResponse implements StatusToXContentOb
     private static final ParseField TIMED_OUT = new ParseField("timed_out");
     private static final ParseField TERMINATED_EARLY = new ParseField("terminated_early");
     private static final ParseField NUM_REDUCE_PHASES = new ParseField("num_reduce_phases");
+    private static final ParseField EXT = new ParseField("ext");
 
     private final SearchResponseSections internalResponse;
     private final String scrollId;
@@ -91,6 +93,8 @@ public class SearchResponse extends ActionResponse implements StatusToXContentOb
     private final Clusters clusters;
     private final long tookInMillis;
 
+    private List<SearchExtBuilder> searchExtBuilders = new ArrayList<>();
+
     public SearchResponse(StreamInput in) throws IOException {
         super(in);
         internalResponse = new InternalSearchResponse(in);
@@ -312,6 +316,15 @@ public XContentBuilder innerToXContent(XContentBuilder builder, Params params) t
         );
         clusters.toXContent(builder, params);
         internalResponse.toXContent(builder, params);
+
+        if (!searchExtBuilders.isEmpty()) {
+            builder.startObject(EXT.getPreferredName());
+            for (SearchExtBuilder searchExtBuilder : searchExtBuilders) {
+                searchExtBuilder.toXContent(builder, params);
+            }
+            builder.endObject();
+        }
+
         return builder;
     }
 
@@ -466,6 +479,14 @@ public String toString() {
         return Strings.toString(MediaTypeRegistry.JSON, this);
     }
 
+    public void addSearchExtBuilder(SearchExtBuilder searchExtBuilder) {
+        this.searchExtBuilders.add(searchExtBuilder);
+    }
+
+    public List<SearchExtBuilder> getSearchExtBuilders() {
+        return this.searchExtBuilders;
+    }
+
     /**
      * Holds info about the clusters that the search was executed on: how many in total, how many of them were successful
      * and how many of them were skipped.

diff --git a/server/src/test/java/org/opensearch/action/search/SearchResponseTests.java b/server/src/test/java/org/opensearch/action/search/SearchResponseTests.java
@@ -40,6 +40,7 @@
 import org.opensearch.core.common.Strings;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
+import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.xcontent.MediaType;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
@@ -48,6 +49,7 @@
 import org.opensearch.core.xcontent.XContentHelper;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.rest.action.search.RestSearchAction;
+import org.opensearch.search.SearchExtBuilder;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
 import org.opensearch.search.SearchHitsTests;
@@ -264,6 +266,7 @@ public void testToXContent() {
                 ShardSearchFailure.EMPTY_ARRAY,
                 SearchResponse.Clusters.EMPTY
             );
+            response.addSearchExtBuilder(new DummySearchExtBuilder());
             StringBuilder expectedString = new StringBuilder();
             expectedString.append("{");
             {
@@ -280,11 +283,17 @@ public void testToXContent() {
                 {
                     expectedString.append("{\"total\":{\"value\":100,\"relation\":\"eq\"},");
                     expectedString.append("\"max_score\":1.5,");
-                    expectedString.append("\"hits\":[{\"_id\":\"id1\",\"_score\":2.0}]}");
+                    expectedString.append("\"hits\":[{\"_id\":\"id1\",\"_score\":2.0}]},");
+                }
+                expectedString.append("\"ext\":");
+                {
+                    expectedString.append("{\"dummy\":\"foo\"}");
                 }
             }
             expectedString.append("}");
             assertEquals(expectedString.toString(), Strings.toString(MediaTypeRegistry.JSON, response));
+            List<SearchExtBuilder> searchExtBuilders = response.getSearchExtBuilders();
+            assertEquals(1, searchExtBuilders.size());
         }
         {
             SearchResponse response = new SearchResponse(
@@ -368,4 +377,32 @@ public void testToXContentEmptyClusters() throws IOException {
         deserialized.getClusters().toXContent(builder, ToXContent.EMPTY_PARAMS);
         assertEquals(0, builder.toString().length());
     }
+
+    static class DummySearchExtBuilder extends SearchExtBuilder {
+
+        @Override
+        public String getWriteableName() {
+            return "DummySearchExtBuilder";
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            out.writeString("foo");
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            return builder.field("dummy", "foo");
+        }
+
+        @Override
+        public int hashCode() {
+            return 0;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            return false;
+        }
+    }
 }