Added SearchTemplate support for reactive client

bilak · web-flow · commit 005d6a4d6ff1 · 2021-03-16T21:37:55.000+01:00
Original Pull Request #1726 Closes #1725
diff --git a/src/main/java/org/springframework/data/elasticsearch/client/reactive/DefaultReactiveElasticsearchClient.java b/src/main/java/org/springframework/data/elasticsearch/client/reactive/DefaultReactiveElasticsearchClient.java
@@ -92,6 +92,8 @@
 import org.elasticsearch.index.reindex.UpdateByQueryRequest;
 import org.elasticsearch.rest.BytesRestResponse;
 import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.script.mustache.SearchTemplateRequest;
+import org.elasticsearch.script.mustache.SearchTemplateResponse;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchHits;
 import org.elasticsearch.search.aggregations.Aggregation;
@@ -404,6 +406,13 @@ public Mono<Long> count(HttpHeaders headers, SearchRequest searchRequest) {
 				.next();
 	}
 
+	@Override
+	public Flux<SearchHit> searchTemplate(HttpHeaders headers, SearchTemplateRequest searchTemplateRequest) {
+		return sendRequest(searchTemplateRequest, requestCreator.searchTemplate(), SearchTemplateResponse.class, headers)
+			.map(r ->  r.getResponse().getHits())
+			.flatMap(Flux::fromIterable);
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.elasticsearch.client.reactive.ReactiveElasticsearchClient#ping(org.springframework.http.HttpHeaders, org.elasticsearch.action.search.SearchRequest)
diff --git a/src/main/java/org/springframework/data/elasticsearch/client/reactive/ReactiveElasticsearchClient.java b/src/main/java/org/springframework/data/elasticsearch/client/reactive/ReactiveElasticsearchClient.java
@@ -52,6 +52,7 @@
 import org.elasticsearch.index.reindex.BulkByScrollResponse;
 import org.elasticsearch.index.reindex.DeleteByQueryRequest;
 import org.elasticsearch.index.reindex.UpdateByQueryRequest;
+import org.elasticsearch.script.mustache.SearchTemplateRequest;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.aggregations.Aggregation;
 import org.elasticsearch.search.suggest.Suggest;
@@ -385,6 +386,43 @@ default Mono<Long> count(SearchRequest searchRequest) {
 	 */
 	Mono<Long> count(HttpHeaders headers, SearchRequest searchRequest);
 
+	/**
+	 * Executes a {@link SearchTemplateRequest} against the {@literal search template} API.
+	 *
+	 * @param consumer must not be {@literal null}.
+	 * @see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/search-template.html">Search Template
+	 * 			API on elastic.co</a>
+	 * @return the {@link Flux} emitting {@link SearchHit hits} one by one.
+	 */
+	default Flux<SearchHit> searchTemplate(Consumer<SearchTemplateRequest> consumer) {
+		SearchTemplateRequest request = new SearchTemplateRequest();
+		consumer.accept(request);
+		return searchTemplate(request);
+	}
+
+	/**
+	 * Executes a {@link SearchTemplateRequest} against the {@literal search template} API.
+	 *
+	 * @param searchTemplateRequest must not be {@literal null}.
+	 * @see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/search-template.html">Search Template
+	 * 			API on elastic.co</a>
+	 * @return the {@link Flux} emitting {@link SearchHit hits} one by one.
+	 */
+	default Flux<SearchHit> searchTemplate(SearchTemplateRequest searchTemplateRequest) {
+		return searchTemplate(HttpHeaders.EMPTY, searchTemplateRequest);
+	}
+
+	/**
+	 * Executes a {@link SearchTemplateRequest} against the {@literal search template} API.
+	 *
+	 * @param headers Use {@link HttpHeaders} to provide eg. authentication data. Must not be {@literal null}.
+	 * @param searchTemplateRequest must not be {@literal null}.
+	 * @see <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/search-template.html">Search Template
+	 * 			API on elastic.co</a>
+	 * @return the {@link Flux} emitting {@link SearchHit hits} one by one.
+	 */
+	Flux<SearchHit> searchTemplate(HttpHeaders headers, SearchTemplateRequest searchTemplateRequest);
+
 	/**
 	 * Execute a {@link SearchRequest} against the {@literal search} API.
 	 *
diff --git a/src/main/java/org/springframework/data/elasticsearch/client/reactive/RequestCreator.java b/src/main/java/org/springframework/data/elasticsearch/client/reactive/RequestCreator.java
@@ -34,6 +34,7 @@
 import org.elasticsearch.client.indices.PutMappingRequest;
 import org.elasticsearch.index.reindex.DeleteByQueryRequest;
 import org.elasticsearch.index.reindex.UpdateByQueryRequest;
+import org.elasticsearch.script.mustache.SearchTemplateRequest;
 import org.springframework.data.elasticsearch.UncategorizedElasticsearchException;
 import org.springframework.data.elasticsearch.client.util.RequestConverters;
 
@@ -49,6 +50,10 @@ default Function<SearchRequest, Request> search() {
 		return RequestConverters::search;
 	}
 
+	default Function<SearchTemplateRequest, Request> searchTemplate() {
+		return RequestConverters::searchTemplate;
+	}
+
 	default Function<SearchScrollRequest, Request> scroll() {
 		return RequestConverters::searchScroll;
 	}
diff --git a/src/main/java/org/springframework/data/elasticsearch/client/util/RequestConverters.java b/src/main/java/org/springframework/data/elasticsearch/client/util/RequestConverters.java
@@ -98,6 +98,7 @@
 import org.elasticsearch.index.reindex.ReindexRequest;
 import org.elasticsearch.index.reindex.UpdateByQueryRequest;
 import org.elasticsearch.index.seqno.SequenceNumbers;
+import org.elasticsearch.script.mustache.SearchTemplateRequest;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.tasks.TaskId;
 import org.springframework.data.elasticsearch.client.reactive.ReactiveElasticsearchClient;
@@ -411,6 +412,21 @@ public static Request search(SearchRequest searchRequest) {
 		return request;
 	}
 
+	public static Request searchTemplate(SearchTemplateRequest templateRequest) {
+		SearchRequest searchRequest = templateRequest.getRequest();
+
+		String endpoint = new EndpointBuilder().addCommaSeparatedPathParts(templateRequest.getRequest().indices())
+			.addPathPart("_search").addPathPart("template").build();
+
+		Request request = new Request(HttpMethod.GET.name(), endpoint);
+		Params params = new Params(request);
+		addSearchRequestParams(params, searchRequest);
+
+		request.setEntity(createEntity(templateRequest, REQUEST_BODY_CONTENT_TYPE));
+
+		return request;
+	}
+
 	/**
 	 * Creates a count request.
 	 *
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/SearchOperations.java b/src/main/java/org/springframework/data/elasticsearch/core/SearchOperations.java
@@ -68,7 +68,7 @@ default long count(Query query, IndexCoordinates index) {
 	 * Does a suggest query
 	 *
 	 * @param suggestion the query
-	 * @param the entity class
+	 * @param clazz the entity class
 	 * @return the suggest response
 	 * @since 4.1
 	 */
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java b/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java
@@ -20,6 +20,7 @@
 import java.util.List;
 
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.script.mustache.SearchTemplateRequestBuilder;
 import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.collapse.CollapseBuilder;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
@@ -50,6 +51,7 @@ public class NativeSearchQuery extends AbstractQuery {
 	@Nullable private HighlightBuilder highlightBuilder;
 	@Nullable private HighlightBuilder.Field[] highlightFields;
 	@Nullable private List<IndexBoost> indicesBoost;
+	@Nullable private SearchTemplateRequestBuilder searchTemplate;
 
 	public NativeSearchQuery(@Nullable QueryBuilder query) {
 
@@ -163,4 +165,12 @@ public void setIndicesBoost(List<IndexBoost> indicesBoost) {
 		this.indicesBoost = indicesBoost;
 	}
 
+	@Nullable
+	public SearchTemplateRequestBuilder getSearchTemplate() {
+		return searchTemplate;
+	}
+
+	public void setSearchTemplate(@Nullable SearchTemplateRequestBuilder searchTemplate) {
+		this.searchTemplate = searchTemplate;
+	}
 }
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java b/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java
@@ -25,6 +25,7 @@
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.script.mustache.SearchTemplateRequestBuilder;
 import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.collapse.CollapseBuilder;
 import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
@@ -61,6 +62,7 @@ public class NativeSearchQueryBuilder {
 	@Nullable private SourceFilter sourceFilter;
 	@Nullable private CollapseBuilder collapseBuilder;
 	@Nullable private List<IndexBoost> indicesBoost;
+	@Nullable private SearchTemplateRequestBuilder searchTemplateBuilder;
 	private float minScore;
 	private boolean trackScores;
 	@Nullable private Collection<String> ids;
@@ -118,6 +120,11 @@ public NativeSearchQueryBuilder withIndicesBoost(List<IndexBoost> indicesBoost)
 		return this;
 	}
 
+	public NativeSearchQueryBuilder withSearchTemplate(SearchTemplateRequestBuilder searchTemplateBuilder){
+		this.searchTemplateBuilder = searchTemplateBuilder;
+		return this;
+	}
+
 	public NativeSearchQueryBuilder withPageable(Pageable pageable) {
 		this.pageable = pageable;
 		return this;
@@ -216,6 +223,10 @@ public NativeSearchQuery build() {
 			nativeSearchQuery.setIndicesBoost(indicesBoost);
 		}
 
+		if (searchTemplateBuilder != null) {
+			nativeSearchQuery.setSearchTemplate(searchTemplateBuilder);
+		}
+
 		if (!isEmpty(scriptFields)) {
 			nativeSearchQuery.setScriptFields(scriptFields);
 		}
diff --git a/src/test/java/org/springframework/data/elasticsearch/client/reactive/ReactiveElasticsearchClientIntegrationTests.java b/src/test/java/org/springframework/data/elasticsearch/client/reactive/ReactiveElasticsearchClientIntegrationTests.java
@@ -52,6 +52,7 @@
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptType;
+import org.elasticsearch.script.mustache.SearchTemplateRequest;
 import org.elasticsearch.search.aggregations.AggregationBuilders;
 import org.elasticsearch.search.aggregations.bucket.terms.StringTerms;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
@@ -73,6 +74,7 @@
 import org.springframework.data.elasticsearch.junit.jupiter.ReactiveElasticsearchRestTemplateConfiguration;
 import org.springframework.data.elasticsearch.junit.jupiter.SpringIntegrationTest;
 import org.springframework.http.HttpHeaders;
+import org.springframework.http.MediaType;
 import org.springframework.test.context.ContextConfiguration;
 
 /**
@@ -434,6 +436,63 @@ public void deleteShouldReturnNotFoundForNonExistingDocument() {
 				.verifyComplete();
 	}
 
+	@Test // #1725
+	public void inlineSearchTemplateShouldFindMatchingDocuments() {
+
+		addSourceDocument().to(INDEX_I);
+		addSourceDocument().to(INDEX_I);
+
+		Map<String, Object> testDoc = new LinkedHashMap<>();
+		testDoc.put("firstname", "inline");
+		testDoc.put("lastname", "template");
+		add(testDoc).to(INDEX_I);
+
+		SearchTemplateRequest request = new SearchTemplateRequest(new SearchRequest(INDEX_I));
+		request.setScriptType(ScriptType.INLINE);
+		request.setScript("{\"query\":{\"match\":{\"firstname\":\"{{firstname}}\"}}}");
+		Map<String, Object> params = new LinkedHashMap<>();
+		params.put("firstname", "inline");
+		request.setScriptParams(params);
+
+		client.searchTemplate(request)
+			.as(StepVerifier::create)
+			.expectNextCount(1)
+			.verifyComplete();
+	}
+
+	@Test // #1725
+	public void storedSearchTemplateShouldFindMatchingDocuments() {
+
+		addSourceDocument().to(INDEX_I);
+		addSourceDocument().to(INDEX_I);
+
+		Map<String, Object> testDoc = new LinkedHashMap<>();
+		testDoc.put("firstname", "stored");
+		testDoc.put("lastname", "template");
+		add(testDoc).to(INDEX_I);
+
+		client.execute(c -> c.post()
+			.uri(builder -> builder.path("_scripts/searchbyfirstname").build())
+			.contentType(MediaType.APPLICATION_JSON)
+			.bodyValue(
+				"{\"script\":{\"lang\":\"mustache\",\"source\":{\"query\":{\"match\":{\"firstname\":\"{{firstname}}\"}}}}}")
+			.retrieve()
+			.bodyToMono(Void.class))
+			.block();
+
+		SearchTemplateRequest request = new SearchTemplateRequest(new SearchRequest(INDEX_I));
+		request.setScriptType(ScriptType.STORED);
+		request.setScript("searchbyfirstname");
+		Map<String, Object> params = new LinkedHashMap<>();
+		params.put("firstname", "stored");
+		request.setScriptParams(params);
+
+		client.searchTemplate(request)
+			.as(StepVerifier::create)
+			.expectNextCount(1)
+			.verifyComplete();
+	}
+
 	@Test // DATAES-488
 	public void searchShouldFindExistingDocuments() {
 

Original file line number	Diff line number	Diff line change
`@@ -68,7 +68,7 @@ default long count(Query query, IndexCoordinates index) {`
`68`	`68`	`* Does a suggest query`
`69`	`69`	`*`
`70`	`70`	`* @param suggestion the query`
`71`		`- * @param the entity class`
	`71`	`+ * @param clazz the entity class`
`72`	`72`	`* @return the suggest response`
`73`	`73`	`* @since 4.1`
`74`	`74`	`*/`