Skip to content

Commit

Permalink
Fix _source bug (#749)
Browse files Browse the repository at this point in the history 
Signed-off-by: Amit Galitzky <[email protected]>
  • Loading branch information
@amitgalitz
amitgalitz authored Dec 7, 2022
1 parent 867c1b3 commit dc293ac
Show file tree
Hide file tree
Showing 5 changed files with 61 additions and 56 deletions.
2 changes: 1 addition & 1 deletion src/main/java/org/opensearch/ad/model/AnomalyDetector.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ public class AnomalyDetector implements Writeable, ToXContentObject {
* @param detectionInterval detecting interval
* @param windowDelay max delay window for realtime data
* @param shingleSize number of the most recent time intervals to form a shingled data point
* @param uiMetadata metadata used by Kibana
* @param uiMetadata metadata used by OpenSearch-Dashboards
* @param schemaVersion anomaly detector index mapping version
* @param lastUpdateTime detector's last update time
* @param categoryFields a list of partition fields
Expand Down
44 changes: 4 additions & 40 deletions src/main/java/org/opensearch/ad/rest/AbstractSearchAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,6 @@

import static org.opensearch.ad.util.RestHandlerUtils.getSourceContext;
import static org.opensearch.common.xcontent.ToXContent.EMPTY_PARAMS;
import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
import static org.opensearch.common.xcontent.XContentParserUtils.ensureExpectedToken;

import java.io.IOException;
import java.util.ArrayList;
Expand All @@ -27,25 +25,16 @@
import org.opensearch.action.search.SearchRequest;
import org.opensearch.action.search.SearchResponse;
import org.opensearch.ad.constant.CommonErrorMessages;
import org.opensearch.ad.model.AnomalyDetector;
import org.opensearch.ad.rest.handler.AnomalyDetectorFunction;
import org.opensearch.ad.settings.EnabledSetting;
import org.opensearch.client.node.NodeClient;
import org.opensearch.common.bytes.BytesReference;
import org.opensearch.common.util.concurrent.ThreadContext;
import org.opensearch.common.xcontent.LoggingDeprecationHandler;
import org.opensearch.common.xcontent.ToXContentObject;
import org.opensearch.common.xcontent.XContentBuilder;
import org.opensearch.common.xcontent.XContentParser;
import org.opensearch.common.xcontent.XContentType;
import org.opensearch.rest.BaseRestHandler;
import org.opensearch.rest.BytesRestResponse;
import org.opensearch.rest.RestChannel;
import org.opensearch.rest.RestRequest;
import org.opensearch.rest.RestResponse;
import org.opensearch.rest.RestStatus;
import org.opensearch.rest.action.RestResponseListener;
import org.opensearch.search.SearchHit;
import org.opensearch.search.builder.SearchSourceBuilder;

/**
Expand Down Expand Up @@ -82,21 +71,15 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
}
SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
searchSourceBuilder.parseXContent(request.contentOrSourceParamParser());
searchSourceBuilder.fetchSource(getSourceContext(request));
// order of response will be re-arranged everytime we use `_source`, we sometimes do this
// even if user doesn't give this field as we exclude ui_metadata if request isn't from OSD
// ref-link: https://github.com/elastic/elasticsearch/issues/17639
searchSourceBuilder.fetchSource(getSourceContext(request, searchSourceBuilder));
searchSourceBuilder.seqNoAndPrimaryTerm(true).version(true);
SearchRequest searchRequest = new SearchRequest().source(searchSourceBuilder).indices(this.index);
return channel -> client.execute(actionType, searchRequest, search(channel));
}

protected void executeWithAdmin(NodeClient client, AnomalyDetectorFunction function, RestChannel channel) {
try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
function.execute();
} catch (Exception e) {
logger.error("Failed to execute with admin", e);
onFailure(channel, e);
}
}

protected void onFailure(RestChannel channel, Exception e) {
try {
channel.sendResponse(new BytesRestResponse(channel, e));
Expand All @@ -112,25 +95,6 @@ public RestResponse buildResponse(SearchResponse response) throws Exception {
if (response.isTimedOut()) {
return new BytesRestResponse(RestStatus.REQUEST_TIMEOUT, response.toString());
}

if (clazz == AnomalyDetector.class) {
for (SearchHit hit : response.getHits()) {
XContentParser parser = XContentType.JSON
.xContent()
.createParser(
channel.request().getXContentRegistry(),
LoggingDeprecationHandler.INSTANCE,
hit.getSourceAsString()
);
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);

// write back id and version to anomaly detector object
ToXContentObject xContentObject = AnomalyDetector.parse(parser, hit.getId(), hit.getVersion());
XContentBuilder builder = xContentObject.toXContent(jsonBuilder(), EMPTY_PARAMS);
hit.sourceRef(BytesReference.bytes(builder));
}
}

return new BytesRestResponse(RestStatus.OK, response.toXContent(channel.newBuilder(), EMPTY_PARAMS));
}
};
Expand Down
2 changes: 1 addition & 1 deletion src/main/java/org/opensearch/ad/rest/RestSearchAnomalyResultAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli

SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
searchSourceBuilder.parseXContent(request.contentOrSourceParamParser());
searchSourceBuilder.fetchSource(getSourceContext(request));
searchSourceBuilder.fetchSource(getSourceContext(request, searchSourceBuilder));
searchSourceBuilder.seqNoAndPrimaryTerm(true).version(true);
SearchRequest searchRequest = new SearchRequest().source(searchSourceBuilder).indices(this.index);

Expand Down
29 changes: 23 additions & 6 deletions src/main/java/org/opensearch/ad/util/RestHandlerUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import java.util.List;
import java.util.Set;

import org.apache.commons.lang.ArrayUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.OpenSearchStatusException;
Expand All @@ -43,6 +44,7 @@
import org.opensearch.rest.RestChannel;
import org.opensearch.rest.RestRequest;
import org.opensearch.rest.RestStatus;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.fetch.subphase.FetchSourceContext;

import com.google.common.base.Throwables;
Expand Down Expand Up @@ -80,21 +82,36 @@ public final class RestHandlerUtils {
public static final String VALIDATE = "_validate";
public static final ToXContent.MapParams XCONTENT_WITH_TYPE = new ToXContent.MapParams(ImmutableMap.of("with_type", "true"));

private static final String OPENSEARCH_DASHBOARDS_USER_AGENT = "OpenSearch Dashboards";
private static final String[] UI_METADATA_EXCLUDE = new String[] { AnomalyDetector.UI_METADATA_FIELD };
public static final String OPENSEARCH_DASHBOARDS_USER_AGENT = "OpenSearch Dashboards";
public static final String[] UI_METADATA_EXCLUDE = new String[] { AnomalyDetector.UI_METADATA_FIELD };

private RestHandlerUtils() {}

/**
* Checks to see if the request came from Kibana, if so we want to return the UI Metadata from the document.
* Checks to see if the request came from OpenSearch-Dashboards, if so we want to return the UI Metadata from the document.
* If the request came from the client then we exclude the UI Metadata from the search result.
*
* We also take into account the given `_source` field and respect the correct fields to be returned.
* @param request rest request
* @param searchSourceBuilder an instance of the searchSourceBuilder to fetch _source field
* @return instance of {@link org.opensearch.search.fetch.subphase.FetchSourceContext}
*/
public static FetchSourceContext getSourceContext(RestRequest request) {
public static FetchSourceContext getSourceContext(RestRequest request, SearchSourceBuilder searchSourceBuilder) {
String userAgent = Strings.coalesceToEmpty(request.header("User-Agent"));
if (!userAgent.contains(OPENSEARCH_DASHBOARDS_USER_AGENT)) {

// If there is a _source given in request than we either add UI_Metadata to exclude or not depending on if request
// is from OpenSearch-Dashboards, if no _source field then we either exclude UI_metadata or return nothing at all.
if (searchSourceBuilder.fetchSource() != null) {
if (userAgent.contains(OPENSEARCH_DASHBOARDS_USER_AGENT)) {
return new FetchSourceContext(
true,
searchSourceBuilder.fetchSource().includes(),
searchSourceBuilder.fetchSource().excludes()
);
} else {
String[] newArray = (String[]) ArrayUtils.addAll(searchSourceBuilder.fetchSource().excludes(), UI_METADATA_EXCLUDE);
return new FetchSourceContext(true, searchSourceBuilder.fetchSource().includes(), newArray);
}
} else if (!userAgent.contains(OPENSEARCH_DASHBOARDS_USER_AGENT)) {
return new FetchSourceContext(true, Strings.EMPTY_ARRAY, UI_METADATA_EXCLUDE);
} else {
return null;
Expand Down
40 changes: 32 additions & 8 deletions src/test/java/org/opensearch/ad/util/RestHandlerUtilsTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

import static org.opensearch.ad.TestHelpers.builder;
import static org.opensearch.ad.TestHelpers.randomFeature;
import static org.opensearch.ad.util.RestHandlerUtils.OPENSEARCH_DASHBOARDS_USER_AGENT;

import java.io.IOException;

Expand All @@ -24,6 +25,7 @@
import org.opensearch.common.xcontent.XContentParser;
import org.opensearch.rest.RestChannel;
import org.opensearch.rest.RestRequest;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.fetch.subphase.FetchSourceContext;
import org.opensearch.test.OpenSearchTestCase;
import org.opensearch.test.rest.FakeRestChannel;
Expand All @@ -34,17 +36,39 @@

public class RestHandlerUtilsTests extends OpenSearchTestCase {

public void testGetSourceContext() {
RestRequest request = new FakeRestRequest();
FetchSourceContext context = RestHandlerUtils.getSourceContext(request);
assertArrayEquals(new String[] { "ui_metadata" }, context.excludes());
public void testGetSourceContextFromOpenSearchDashboardEmptyExcludes() {
FakeRestRequest.Builder builder = new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY);
builder.withHeaders(ImmutableMap.of("User-Agent", ImmutableList.of(OPENSEARCH_DASHBOARDS_USER_AGENT, randomAlphaOfLength(10))));
SearchSourceBuilder testSearchSourceBuilder = new SearchSourceBuilder();
testSearchSourceBuilder.fetchSource(new String[] { "a" }, new String[0]);
FetchSourceContext sourceContext = RestHandlerUtils.getSourceContext(builder.build(), testSearchSourceBuilder);
assertArrayEquals(new String[] { "a" }, sourceContext.includes());
assertEquals(0, sourceContext.excludes().length);
assertEquals(1, sourceContext.includes().length);
}

public void testGetSourceContextFromClientWithExcludes() {
FakeRestRequest.Builder builder = new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY);
SearchSourceBuilder testSearchSourceBuilder = new SearchSourceBuilder();
testSearchSourceBuilder.fetchSource(new String[] { "a" }, new String[] { "b" });
FetchSourceContext sourceContext = RestHandlerUtils.getSourceContext(builder.build(), testSearchSourceBuilder);
assertEquals(sourceContext.excludes().length, 2);
}

public void testGetSourceContextFromClientWithoutSource() {
FakeRestRequest.Builder builder = new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY);
SearchSourceBuilder testSearchSourceBuilder = new SearchSourceBuilder();
FetchSourceContext sourceContext = RestHandlerUtils.getSourceContext(builder.build(), testSearchSourceBuilder);
assertEquals(sourceContext.excludes().length, 1);
assertEquals(sourceContext.includes().length, 0);
}

public void testGetSourceContextForKibana() {
public void testGetSourceContextOpenSearchDashboardWithoutSources() {
FakeRestRequest.Builder builder = new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY);
builder.withHeaders(ImmutableMap.of("User-Agent", ImmutableList.of("OpenSearch Dashboards", randomAlphaOfLength(10))));
FetchSourceContext context = RestHandlerUtils.getSourceContext(builder.build());
assertNull(context);
builder.withHeaders(ImmutableMap.of("User-Agent", ImmutableList.of(OPENSEARCH_DASHBOARDS_USER_AGENT, randomAlphaOfLength(10))));
SearchSourceBuilder testSearchSourceBuilder = new SearchSourceBuilder();
FetchSourceContext sourceContext = RestHandlerUtils.getSourceContext(builder.build(), testSearchSourceBuilder);
assertNull(sourceContext);
}

public void testCreateXContentParser() throws IOException {
Expand Down

0 comments on commit dc293ac

Please sign in to comment.