Skip to content

Commit

Permalink
Add MultiSearchTemplate support to High Level Rest client
Browse files Browse the repository at this point in the history
Note some changes to core classes were required to support tests:

* Added SearchType.currentlySupported to list only the usable options
* Changed SearchResponse.Clusters constructor to public
* Added SearchRequest.DEFAULT_BATCHED_REDUCE_SIZE constant
* Added missing "tookInMillis" to MultiSearchTemplateResponse
* Added SearchTemplateResponse toString() for comparisons in tests
  • Loading branch information
markharwood committed Jun 12, 2018
1 parent b2e48c9 commit 4f8451f
Show file tree
Hide file tree
Showing 17 changed files with 734 additions and 84 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@
import org.elasticsearch.index.VersionType;
import org.elasticsearch.index.rankeval.RankEvalRequest;
import org.elasticsearch.rest.action.search.RestSearchAction;
import org.elasticsearch.script.mustache.MultiSearchTemplateRequest;
import org.elasticsearch.script.mustache.SearchTemplateRequest;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.tasks.TaskId;
Expand Down Expand Up @@ -573,6 +574,21 @@ static Request searchTemplate(SearchTemplateRequest searchTemplateRequest) throw
request.setEntity(createEntity(searchTemplateRequest, REQUEST_BODY_CONTENT_TYPE));
return request;
}

static Request multiSearchTemplate(MultiSearchTemplateRequest multiSearchTemplateRequest) throws IOException {
Request request = new Request(HttpPost.METHOD_NAME, "/_msearch/template");

Params params = new Params(request);
params.putParam(RestSearchAction.TYPED_KEYS_PARAM, "true");
if (multiSearchTemplateRequest.maxConcurrentSearchRequests() != MultiSearchRequest.MAX_CONCURRENT_SEARCH_REQUESTS_DEFAULT) {
params.putParam("max_concurrent_searches", Integer.toString(multiSearchTemplateRequest.maxConcurrentSearchRequests()));
}

XContent xContent = REQUEST_BODY_CONTENT_TYPE.xContent();
byte[] source = MultiSearchTemplateRequest.writeMultiLineFormat(multiSearchTemplateRequest, xContent);
request.setEntity(new ByteArrayEntity(source, createContentType(xContent.type())));
return request;
}

static Request existsAlias(GetAliasesRequest getAliasesRequest) {
if ((getAliasesRequest.indices() == null || getAliasesRequest.indices().length == 0) &&
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,8 @@
import org.elasticsearch.plugins.spi.NamedXContentProvider;
import org.elasticsearch.rest.BytesRestResponse;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.script.mustache.MultiSearchTemplateRequest;
import org.elasticsearch.script.mustache.MultiSearchTemplateResponse;
import org.elasticsearch.script.mustache.SearchTemplateRequest;
import org.elasticsearch.script.mustache.SearchTemplateResponse;
import org.elasticsearch.search.aggregations.Aggregation;
Expand Down Expand Up @@ -909,6 +911,32 @@ public final RankEvalResponse rankEval(RankEvalRequest rankEvalRequest, RequestO
emptySet());
}


/**
* Executes a request using the Multi Search Template API.
*
* See <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/multi-search-template.html">Multi Search Template API
* on elastic.co</a>.
*/
public final MultiSearchTemplateResponse multiSearchTemplate(MultiSearchTemplateRequest multiSearchTemplateRequest,
RequestOptions options) throws IOException {
return performRequestAndParseEntity(multiSearchTemplateRequest, RequestConverters::multiSearchTemplate,
options, MultiSearchTemplateResponse::fromXContext, emptySet());
}

/**
* Asynchronously executes a request using the Multi Search Template API
*
* See <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/multi-search-template.html">Multi Search Template API
* on elastic.co</a>.
*/
public final void multiSearchTemplateAsync(MultiSearchTemplateRequest multiSearchTemplateRequest,
RequestOptions options,
ActionListener<MultiSearchTemplateResponse> listener) {
performRequestAsyncAndParseEntity(multiSearchTemplateRequest, RequestConverters::multiSearchTemplate,
options, MultiSearchTemplateResponse::fromXContext, listener, emptySet());
}

/**
* Executes a request using the Ranking Evaluation API.
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,12 @@
import org.elasticsearch.index.rankeval.RatedRequest;
import org.elasticsearch.index.rankeval.RestRankEvalAction;
import org.elasticsearch.repositories.fs.FsRepository;
import org.elasticsearch.rest.RestRequest;
import org.elasticsearch.rest.RestRequest.Method;
import org.elasticsearch.rest.action.search.RestSearchAction;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.script.mustache.MultiSearchTemplateRequest;
import org.elasticsearch.script.mustache.RestMultiSearchTemplateAction;
import org.elasticsearch.script.mustache.SearchTemplateRequest;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
Expand Down Expand Up @@ -1304,7 +1308,53 @@ public void testRenderSearchTemplate() throws Exception {
assertEquals(Collections.emptyMap(), request.getParameters());
assertToXContentBody(searchTemplateRequest, request.getEntity());
}


public void testMultiSearchTemplate() throws Exception {
final int numSearchRequests = randomIntBetween(1, 10);
MultiSearchTemplateRequest multiSearchTemplateRequest = new MultiSearchTemplateRequest();

for (int i = 0; i < numSearchRequests; i++) {
// Create a random request.
String[] indices = randomIndicesNames(0, 5);
SearchRequest searchRequest = new SearchRequest(indices);

Map<String, String> expectedParams = new HashMap<>();
setRandomSearchParams(searchRequest, expectedParams);

// scroll is not supported in the current msearch or msearchtemplate api, so unset it:
searchRequest.scroll((Scroll) null);
// batched reduce size is currently not set-able on a per-request basis as it is a query string parameter only
searchRequest.setBatchedReduceSize(SearchRequest.DEFAULT_BATCHED_REDUCE_SIZE);

setRandomIndicesOptions(searchRequest::indicesOptions, searchRequest::indicesOptions, expectedParams);

SearchTemplateRequest searchTemplateRequest = new SearchTemplateRequest(searchRequest);

searchTemplateRequest.setScript("{\"query\": { \"match\" : { \"{{field}}\" : \"{{value}}\" }}}");
searchTemplateRequest.setScriptType(ScriptType.INLINE);
searchTemplateRequest.setProfile(randomBoolean());

Map<String, Object> scriptParams = new HashMap<>();
scriptParams.put("field", "name");
scriptParams.put("value", randomAlphaOfLengthBetween(2, 5));
searchTemplateRequest.setScriptParams(scriptParams);

multiSearchTemplateRequest.add(searchTemplateRequest);
}

Request multiRequest = RequestConverters.multiSearchTemplate(multiSearchTemplateRequest);

assertEquals(HttpPost.METHOD_NAME, multiRequest.getMethod());
assertEquals("/_msearch/template", multiRequest.getEndpoint());
List<SearchTemplateRequest> searchRequests = multiSearchTemplateRequest.requests();
assertEquals(numSearchRequests, searchRequests.size());

HttpEntity actualEntity = multiRequest.getEntity();
byte[] expectedBytes = MultiSearchTemplateRequest.writeMultiLineFormat(multiSearchTemplateRequest, XContentType.JSON.xContent());
assertEquals(XContentType.JSON.mediaTypeWithoutParameters(), actualEntity.getContentType().getValue());
assertEquals(new BytesArray(expectedBytes), new BytesArray(EntityUtils.toByteArray(actualEntity)));
}

public void testExistsAlias() {
GetAliasesRequest getAliasesRequest = new GetAliasesRequest();
String[] indices = randomBoolean() ? null : randomIndicesNames(0, 5);
Expand Down Expand Up @@ -2013,7 +2063,7 @@ private static void setRandomSearchParams(SearchRequest searchRequest,
expectedParams.put("preference", searchRequest.preference());
}
if (randomBoolean()) {
searchRequest.searchType(randomFrom(SearchType.values()));
searchRequest.searchType(randomFrom(SearchType.CURRENTLY_SUPPORTED));
}
expectedParams.put("search_type", searchRequest.searchType().name().toLowerCase(Locale.ROOT));
if (randomBoolean()) {
Expand All @@ -2038,7 +2088,7 @@ private static void setRandomIndicesOptions(Consumer<IndicesOptions> setter, Sup
Map<String, String> expectedParams) {

if (randomBoolean()) {
setter.accept(IndicesOptions.fromOptions(randomBoolean(), randomBoolean(), randomBoolean(), randomBoolean()));
setter.accept(IndicesOptions.fromOptions(randomBoolean(), randomBoolean(), randomBoolean(), randomBoolean()));
}
expectedParams.put("ignore_unavailable", Boolean.toString(getter.get().ignoreUnavailable()));
expectedParams.put("allow_no_indices", Boolean.toString(getter.get().allowNoIndices()));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,9 @@
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.script.mustache.MultiSearchTemplateRequest;
import org.elasticsearch.script.mustache.MultiSearchTemplateResponse;
import org.elasticsearch.script.mustache.MultiSearchTemplateResponse.Item;
import org.elasticsearch.script.mustache.SearchTemplateRequest;
import org.elasticsearch.script.mustache.SearchTemplateResponse;
import org.elasticsearch.search.SearchHit;
Expand Down Expand Up @@ -854,6 +857,105 @@ public void testRenderSearchTemplate() throws IOException {

assertToXContentEquivalent(expectedSource, actualSource, XContentType.JSON);
}


public void testMultiSearchTemplate() throws Exception {
MultiSearchTemplateRequest multiSearchTemplateRequest = new MultiSearchTemplateRequest();

SearchTemplateRequest goodRequest = new SearchTemplateRequest();
goodRequest.setRequest(new SearchRequest("index"));
goodRequest.setScriptType(ScriptType.INLINE);
goodRequest.setScript(
"{" +
" \"query\": {" +
" \"match\": {" +
" \"num\": {{number}}" +
" }" +
" }" +
"}");
Map<String, Object> scriptParams = new HashMap<>();
scriptParams.put("number", 10);
goodRequest.setScriptParams(scriptParams);
goodRequest.setExplain(true);
goodRequest.setProfile(true);
multiSearchTemplateRequest.add(goodRequest);


SearchTemplateRequest badRequest = new SearchTemplateRequest();
badRequest.setRequest(new SearchRequest("index"));
badRequest.setScriptType(ScriptType.INLINE);
badRequest.setScript("{ NOT VALID JSON {{number}} }");
scriptParams = new HashMap<>();
scriptParams.put("number", 10);
badRequest.setScriptParams(scriptParams);

multiSearchTemplateRequest.add(badRequest);

MultiSearchTemplateResponse multiSearchTemplateResponse =
execute(multiSearchTemplateRequest, highLevelClient()::multiSearchTemplate,
highLevelClient()::multiSearchTemplateAsync);

Item[] responses = multiSearchTemplateResponse.getResponses();

assertEquals(2, responses.length);


assertNull(responses[0].getResponse().getSource());
SearchResponse goodResponse =responses[0].getResponse().getResponse();
assertNotNull(goodResponse);
assertThat(responses[0].isFailure(), Matchers.is(false));
assertEquals(1, goodResponse.getHits().totalHits);
assertEquals(1, goodResponse.getHits().getHits().length);
assertThat(goodResponse.getHits().getMaxScore(), greaterThan(0f));
SearchHit hit = goodResponse.getHits().getHits()[0];
assertNotNull(hit.getExplanation());
assertFalse(goodResponse.getProfileResults().isEmpty());


assertNull(responses[0].getResponse().getSource());
assertThat(responses[1].isFailure(), Matchers.is(true));
assertNotNull(responses[1].getFailureMessage());
assertThat(responses[1].getFailureMessage(), containsString("json_parse_exception"));
}

public void testMultiSearchTemplateAllBad() throws Exception {
MultiSearchTemplateRequest multiSearchTemplateRequest = new MultiSearchTemplateRequest();

SearchTemplateRequest badRequest1 = new SearchTemplateRequest();
badRequest1.setRequest(new SearchRequest("index"));
badRequest1.setScriptType(ScriptType.INLINE);
badRequest1.setScript(
"{" +
" \"query\": {" +
" \"match\": {" +
" \"num\": {{number}}" +
" }" +
" }" +
"}");
Map<String, Object> scriptParams = new HashMap<>();
scriptParams.put("number", "BAD NUMBER");
badRequest1.setScriptParams(scriptParams);
multiSearchTemplateRequest.add(badRequest1);


SearchTemplateRequest badRequest2 = new SearchTemplateRequest();
badRequest2.setRequest(new SearchRequest("index"));
badRequest2.setScriptType(ScriptType.INLINE);
badRequest2.setScript("BAD QUERY TEMPLATE");
scriptParams = new HashMap<>();
scriptParams.put("number", "BAD NUMBER");
badRequest2.setScriptParams(scriptParams);

multiSearchTemplateRequest.add(badRequest2);

// The whole HTTP request should fail if no nested search requests are valid
ElasticsearchStatusException exception = expectThrows(ElasticsearchStatusException.class,
() -> execute(multiSearchTemplateRequest, highLevelClient()::multiSearchTemplate,
highLevelClient()::multiSearchTemplateAsync));

assertEquals(RestStatus.BAD_REQUEST, exception.status());
assertThat(exception.getMessage(), containsString("no requests added"));
}

public void testFieldCaps() throws IOException {
FieldCapabilitiesRequest request = new FieldCapabilitiesRequest()
Expand Down
Loading

0 comments on commit 4f8451f

Please sign in to comment.