Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add SearchExtBuilders to SearchResponse #9379

Merged
merged 11 commits into from
Aug 29, 2023
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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<>();
austintlee marked this conversation as resolved.
Show resolved Hide resolved

public SearchResponse(StreamInput in) throws IOException {
super(in);
internalResponse = new InternalSearchResponse(in);
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -264,6 +266,7 @@ public void testToXContent() {
ShardSearchFailure.EMPTY_ARRAY,
SearchResponse.Clusters.EMPTY
);
response.addSearchExtBuilder(new DummySearchExtBuilder());
StringBuilder expectedString = new StringBuilder();
expectedString.append("{");
{
Expand All @@ -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(
Expand Down Expand Up @@ -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;
}
}
}