Add SearchExtBuilders to SearchResponse

CHANGELOG.md

@@ -79,17 +79,18 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 2.x]
 ### Added
 - Add server version as REST response header [#6583](https://github.com/opensearch-project/OpenSearch/issues/6583)
-- Start replication checkpointTimers on primary before segments upload to remote store. ([#8221]()https://github.com/opensearch-project/OpenSearch/pull/8221)
-- [distribution/archives] [Linux] [x64] Provide the variant of the distributions bundled with JRE ([#8195]()https://github.com/opensearch-project/OpenSearch/pull/8195)
+- Start replication checkpointTimers on primary before segments upload to remote store. ([#8221](https://github.com/opensearch-project/OpenSearch/pull/8221))
+- [distribution/archives] [Linux] [x64] Provide the variant of the distributions bundled with JRE ([#8195](https://github.com/opensearch-project/OpenSearch/pull/8195))
 - Add configuration for file cache size to max remote data ratio to prevent oversubscription of file cache ([#8606](https://github.com/opensearch-project/OpenSearch/pull/8606))
-- Disallow compression level to be set for default and best_compression index codecs ([#8737]()https://github.com/opensearch-project/OpenSearch/pull/8737)
+- Disallow compression level to be set for default and best_compression index codecs ([#8737](https://github.com/opensearch-project/OpenSearch/pull/8737))
 - Prioritize replica shard movement during shard relocation ([#8875](https://github.com/opensearch-project/OpenSearch/pull/8875))
-- Introducing Default and Best Compression codecs as their algorithm name ([#9123]()https://github.com/opensearch-project/OpenSearch/pull/9123)
-- Make SearchTemplateRequest implement IndicesRequest.Replaceable ([#9122]()https://github.com/opensearch-project/OpenSearch/pull/9122)
+- Introducing Default and Best Compression codecs as their algorithm name ([#9123](https://github.com/opensearch-project/OpenSearch/pull/9123))
+- Make SearchTemplateRequest implement IndicesRequest.Replaceable ([#9122](https://github.com/opensearch-project/OpenSearch/pull/9122))
 - [BWC and API enforcement] Define the initial set of annotations, their meaning and relations between them ([#9223](https://github.com/opensearch-project/OpenSearch/pull/9223))
 - [Segment Replication] Support realtime reads for GET requests ([#9212](https://github.com/opensearch-project/OpenSearch/pull/9212))
 - [Feature] Expose term frequency in Painless script score context ([#9081](https://github.com/opensearch-project/OpenSearch/pull/9081))
 - Add support for reading partial files to HDFS repository ([#9513](https://github.com/opensearch-project/OpenSearch/issues/9513))
+- Add support for extensions to search responses using SearchExtBuilder ([#9379](https://github.com/opensearch-project/OpenSearch/pull/9379)) 
 
 ### Dependencies
 - Bump `org.apache.logging.log4j:log4j-core` from 2.17.1 to 2.20.0 ([#8307](https://github.com/opensearch-project/OpenSearch/pull/8307))

server/src/main/java/org/opensearch/action/search/SearchResponse.java

@@ -46,9 +46,12 @@
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.core.xcontent.ToXContentFragment;
 import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.core.xcontent.XContentParseException;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.core.xcontent.XContentParser.Token;
 import org.opensearch.rest.action.RestActions;
+import org.opensearch.search.GenericSearchExtBuilder;
+import org.opensearch.search.SearchExtBuilder;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
 import org.opensearch.search.aggregations.Aggregations;
@@ -65,6 +68,7 @@
 import java.util.Objects;
 import java.util.function.Supplier;
 
+import static org.opensearch.action.search.SearchResponseSections.EXT_FIELD;
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 
 /**
@@ -312,6 +316,7 @@ public XContentBuilder innerToXContent(XContentBuilder builder, Params params) t
         );
         clusters.toXContent(builder, params);
         internalResponse.toXContent(builder, params);
+
         return builder;
     }
 
@@ -339,6 +344,7 @@ public static SearchResponse innerFromXContent(XContentParser parser) throws IOE
         String searchContextId = null;
         List<ShardSearchFailure> failures = new ArrayList<>();
         Clusters clusters = Clusters.EMPTY;
+        List<SearchExtBuilder> extBuilders = new ArrayList<>();
         for (Token token = parser.nextToken(); token != Token.END_OBJECT; token = parser.nextToken()) {
             if (token == Token.FIELD_NAME) {
                 currentFieldName = parser.currentName();
@@ -417,6 +423,33 @@ public static SearchResponse innerFromXContent(XContentParser parser) throws IOE
                         }
                     }
                     clusters = new Clusters(total, successful, skipped);
+                } else if (EXT_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
+                    String extSectionName = null;
+                    while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+                        if (token == XContentParser.Token.FIELD_NAME) {
+                            extSectionName = parser.currentName();
+                        } else {
+                            SearchExtBuilder searchExtBuilder;
+                            try {
+                                searchExtBuilder = parser.namedObject(SearchExtBuilder.class, extSectionName, null);
+                                if (!searchExtBuilder.getWriteableName().equals(extSectionName)) {
+                                    throw new IllegalStateException(
+                                        "The parsed ["
+                                            + searchExtBuilder.getClass().getName()
+                                            + "] object has a "
+                                            + "different writeable name compared to the name of the section that it was parsed from: found ["
+                                            + searchExtBuilder.getWriteableName()
+                                            + "] expected ["
+                                            + extSectionName
+                                            + "]"
+                                    );
+                                }
+                            } catch (XContentParseException e) {
+                                searchExtBuilder = GenericSearchExtBuilder.fromXContent(parser);
+                            }
+                            extBuilders.add(searchExtBuilder);
+                        }
+                    }
                 } else {
                     parser.skipChildren();
                 }
@@ -429,7 +462,8 @@ public static SearchResponse innerFromXContent(XContentParser parser) throws IOE
             timedOut,
             terminatedEarly,
             profile,
-            numReducePhases
+            numReducePhases,
+            extBuilders
         );
         return new SearchResponse(
             searchResponseSections,

server/src/main/java/org/opensearch/action/search/SearchResponseSections.java

@@ -32,18 +32,23 @@
 
 package org.opensearch.action.search;
 
+import org.opensearch.core.ParseField;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.xcontent.ToXContentFragment;
 import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.search.SearchExtBuilder;
 import org.opensearch.search.SearchHits;
 import org.opensearch.search.aggregations.Aggregations;
 import org.opensearch.search.profile.ProfileShardResult;
 import org.opensearch.search.profile.SearchProfileShardResults;
 import org.opensearch.search.suggest.Suggest;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * Base class that holds the various sections which a search response is
@@ -57,13 +62,16 @@
  */
 public class SearchResponseSections implements ToXContentFragment {
 
+    public static final ParseField EXT_FIELD = new ParseField("ext");
+
     protected final SearchHits hits;
     protected final Aggregations aggregations;
     protected final Suggest suggest;
     protected final SearchProfileShardResults profileResults;
     protected final boolean timedOut;
     protected final Boolean terminatedEarly;
     protected final int numReducePhases;
+    protected final List<SearchExtBuilder> searchExtBuilders = new ArrayList<>();
 
     public SearchResponseSections(
         SearchHits hits,
@@ -73,6 +81,19 @@ public SearchResponseSections(
         Boolean terminatedEarly,
         SearchProfileShardResults profileResults,
         int numReducePhases
+    ) {
+        this(hits, aggregations, suggest, timedOut, terminatedEarly, profileResults, numReducePhases, Collections.emptyList());
+    }
+
+    public SearchResponseSections(
+        SearchHits hits,
+        Aggregations aggregations,
+        Suggest suggest,
+        boolean timedOut,
+        Boolean terminatedEarly,
+        SearchProfileShardResults profileResults,
+        int numReducePhases,
+        List<SearchExtBuilder> searchExtBuilders
     ) {
         this.hits = hits;
         this.aggregations = aggregations;
@@ -81,6 +102,7 @@ public SearchResponseSections(
         this.timedOut = timedOut;
         this.terminatedEarly = terminatedEarly;
         this.numReducePhases = numReducePhases;
+        this.searchExtBuilders.addAll(Objects.requireNonNull(searchExtBuilders, "searchExtBuilders must not be null"));
     }
 
     public final boolean timedOut() {
@@ -135,9 +157,20 @@ public final XContentBuilder toXContent(XContentBuilder builder, Params params)
         if (profileResults != null) {
             profileResults.toXContent(builder, params);
         }
+        if (!searchExtBuilders.isEmpty()) {
+            builder.startObject(EXT_FIELD.getPreferredName());
+            for (SearchExtBuilder searchExtBuilder : searchExtBuilders) {
+                searchExtBuilder.toXContent(builder, params);
+            }
+            builder.endObject();
+        }
         return builder;
     }
 
+    public List<SearchExtBuilder> getSearchExtBuilders() {
+        return Collections.unmodifiableList(this.searchExtBuilders);
+    }
+
     protected void writeTo(StreamOutput out) throws IOException {
         throw new UnsupportedOperationException();
     }

server/src/main/java/org/opensearch/search/GenericSearchExtBuilder.java

@@ -0,0 +1,161 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.search;
+
+import org.opensearch.core.ParseField;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.core.xcontent.XContentParseException;
+import org.opensearch.core.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+public final class GenericSearchExtBuilder extends SearchExtBuilder {
+
+    public final static ParseField EXT_BUILDER_NAME = new ParseField("generic_ext");
+
+    private final Object genericObj;
+    private final ValueType valueType;
+
+    enum ValueType {
+        SIMPLE(0),
+        MAP(1),
+        LIST(2);
+
+        private final int value;
+
+        ValueType(int value) {
+            this.value = value;
+        }
+
+        public int getValue() {
+            return value;
+        }
+
+        static ValueType fromInt(int value) {
+            switch (value) {
+                case 0:
+                    return SIMPLE;
+                case 1:
+                    return MAP;
+                case 2:
+                    return LIST;
+                default:
+                    throw new IllegalArgumentException("Unsupported value: " + value);
+            }
+        }
+    }
+
+    public GenericSearchExtBuilder(Object genericObj, ValueType valueType) {
+        this.genericObj = genericObj;
+        this.valueType = valueType;
+    }
+
+    public GenericSearchExtBuilder(StreamInput in) throws IOException {
+        valueType = ValueType.fromInt(in.readInt());
+        switch (valueType) {
+            case SIMPLE:
+                genericObj = in.readGenericValue();
+                break;
+            case MAP:
+                genericObj = in.readMap();
+                break;
+            case LIST:
+                genericObj = in.readList(r -> r.readGenericValue());
+                break;
+            default:
+                throw new IllegalStateException("Unable to construct GenericSearchExtBuilder from incoming stream.");
+        }
+    }
+
+    public static GenericSearchExtBuilder fromXContent(XContentParser parser) throws IOException {
+        // Look at the parser's next token.
+        // If it's START_OBJECT, parse as map, if it's START_ARRAY, parse as list, else
+        // parse as simpleVal
+        XContentParser.Token token = parser.currentToken();
+        ValueType valueType;
+        Object genericObj;
+        if (token == XContentParser.Token.START_OBJECT) {
+            genericObj = parser.map();
+            valueType = ValueType.MAP;
+        } else if (token == XContentParser.Token.START_ARRAY) {
+            genericObj = parser.list();
+            valueType = ValueType.LIST;
+        } else if (token.isValue()) {
+            genericObj = parser.objectText(); // .objectBytes() ??
+            valueType = ValueType.SIMPLE;
+        } else {
+            throw new XContentParseException("Unknown token: " + token);
+        }
+
+        return new GenericSearchExtBuilder(genericObj, valueType);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return EXT_BUILDER_NAME.getPreferredName();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeInt(valueType.getValue());
+        switch (valueType) {
+            case SIMPLE:
+                out.writeGenericValue(genericObj);
+                break;
+            case MAP:
+                out.writeMap((Map<String, Object>) genericObj);
+                break;
+            case LIST:
+                out.writeCollection((List<Object>) genericObj, StreamOutput::writeGenericValue);
+                break;
+            default:
+                throw new IllegalStateException("Unknown valueType: " + valueType);
+        }
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        switch (valueType) {
+            case SIMPLE:
+                return builder.field(EXT_BUILDER_NAME.getPreferredName(), genericObj);
+            case MAP:
+                return builder.field(EXT_BUILDER_NAME.getPreferredName(), (Map<String, Object>) genericObj);
+            case LIST:
+                return builder.field(EXT_BUILDER_NAME.getPreferredName(), (List<Object>) genericObj);
+            default:
+                return null;
+        }
+    }
+
+    // We need this for the equals method.
+    Object getValue() {
+        return genericObj;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(this.valueType, this.genericObj);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == null) {
+            return false;
+        }
+        if (!(obj instanceof GenericSearchExtBuilder)) {
+            return false;
+        }
+        return Objects.equals(getValue(), ((GenericSearchExtBuilder) obj).getValue())
+            && Objects.equals(valueType, ((GenericSearchExtBuilder) obj).valueType);
+    }
+}