DATAES-479 - Allow specifying a HighlightBuilder when creating a query.

Original pull request: #217.

src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java

@@ -134,6 +134,7 @@ import org.springframework.util.StringUtils;
  * @author Alen Turkovic
  * @author Sascha Woo
  * @author Ted Liang
+ * @author Jean-Baptiste Nizet
  */
 public class ElasticsearchTemplate implements ElasticsearchOperations, ApplicationContextAware {
 
@@ -905,8 +906,11 @@ public class ElasticsearchTemplate implements ElasticsearchOperations, Applicati
 			}
 		}
 
-		if (searchQuery.getHighlightFields() != null) {
-			HighlightBuilder highlightBuilder = new HighlightBuilder();
+		if (searchQuery.getHighlightFields() != null || searchQuery.getHighlightBuilder() != null) {
+			HighlightBuilder highlightBuilder = searchQuery.getHighlightBuilder();
+			if (highlightBuilder == null) {
+				highlightBuilder = new HighlightBuilder();
+			}
 			for (HighlightBuilder.Field highlightField : searchQuery.getHighlightFields()) {
 				highlightBuilder.field(highlightField);
 			}
@@ -934,11 +938,11 @@ public class ElasticsearchTemplate implements ElasticsearchOperations, Applicati
 	}
 
 	private SearchResponse getSearchResponse(SearchRequestBuilder requestBuilder) {
-		
+
 		if (QUERY_LOGGER.isDebugEnabled()) {
 			QUERY_LOGGER.debug(requestBuilder.toString());
 		}
-		
+
 		return getSearchResponse(requestBuilder.execute());
 	}
 

src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2017 the original author or authors.
+ * Copyright 2013-2018 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -32,6 +32,7 @@ import java.util.Arrays;
  * @author Rizwan Idrees
  * @author Mohsin Husen
  * @author Artur Konczak
+ * @author Jean-Baptiste Nizet
  */
 public class NativeSearchQuery extends AbstractQuery implements SearchQuery {
 
@@ -41,6 +42,7 @@ public class NativeSearchQuery extends AbstractQuery implements SearchQuery {
     private final List<ScriptField> scriptFields = new ArrayList<>();
 	private List<FacetRequest> facets;
 	private List<AbstractAggregationBuilder> aggregations;
+	private HighlightBuilder highlightBuilder;
 	private HighlightBuilder.Field[] highlightFields;
 	private List<IndexBoost> indicesBoost;
 
@@ -67,6 +69,15 @@ public class NativeSearchQuery extends AbstractQuery implements SearchQuery {
 		this.highlightFields = highlightFields;
 	}
 
+	public NativeSearchQuery(QueryBuilder query, QueryBuilder filter, List<SortBuilder> sorts,
+			HighlightBuilder highlighBuilder, HighlightBuilder.Field[] highlightFields) {
+		this.query = query;
+		this.filter = filter;
+		this.sorts = sorts;
+		this.highlightBuilder = highlighBuilder;
+		this.highlightFields = highlightFields;
+	}
+
 	public QueryBuilder getQuery() {
 		return query;
 	}
@@ -79,6 +90,11 @@ public class NativeSearchQuery extends AbstractQuery implements SearchQuery {
 		return sorts;
 	}
 
+	@Override
+	public HighlightBuilder getHighlightBuilder() {
+		return highlightBuilder;
+	}
+
 	@Override
 	public HighlightBuilder.Field[] getHighlightFields() {
 		return highlightFields;

src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQueryBuilder.java

@@ -37,6 +37,7 @@ import org.springframework.data.elasticsearch.core.facet.FacetRequest;
  * @author Mark Paluch
  * @author Alen Turkovic
  * @author Sascha Woo
+ * @author Jean-Baptiste Nizet
  */
 public class NativeSearchQueryBuilder {
 
@@ -46,6 +47,7 @@ public class NativeSearchQueryBuilder {
 	private List<SortBuilder> sortBuilders = new ArrayList<>();
 	private List<FacetRequest> facetRequests = new ArrayList<>();
 	private List<AbstractAggregationBuilder> aggregationBuilders = new ArrayList<>();
+	private HighlightBuilder highlightBuilder;
 	private HighlightBuilder.Field[] highlightFields;
 	private Pageable pageable = Pageable.unpaged();
 	private String[] indices;
@@ -90,6 +92,11 @@ public class NativeSearchQueryBuilder {
 		return this;
 	}
 
+	public NativeSearchQueryBuilder withHighlightBuilder(HighlightBuilder highlightBuilder) {
+		this.highlightBuilder = highlightBuilder;
+		return this;
+	}
+
 	public NativeSearchQueryBuilder withHighlightFields(HighlightBuilder.Field... highlightFields) {
 		this.highlightFields = highlightFields;
 		return this;
@@ -162,7 +169,7 @@ public class NativeSearchQueryBuilder {
 
 	public NativeSearchQuery build() {
 		NativeSearchQuery nativeSearchQuery = new NativeSearchQuery(queryBuilder, filterBuilder, sortBuilders,
-				highlightFields);
+				highlightBuilder, highlightFields);
 
 		nativeSearchQuery.setPageable(pageable);
 		nativeSearchQuery.setTrackScores(trackScores);

src/main/java/org/springframework/data/elasticsearch/core/query/SearchQuery.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2016 the original author or authors.
+ * Copyright 2013-2018 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -29,6 +29,7 @@ import org.springframework.data.elasticsearch.core.facet.FacetRequest;
  * @author Rizwan Idrees
  * @author Mohsin Husen
  * @author Artur Konczak
+ * @author Jean-Baptiste Nizet
  */
 public interface SearchQuery extends Query {
 
@@ -43,6 +44,8 @@ public interface SearchQuery extends Query {
 
 	List<AbstractAggregationBuilder> getAggregations();
 
+	HighlightBuilder getHighlightBuilder();
+
 	HighlightBuilder.Field[] getHighlightFields();
 
 	List<IndexBoost> getIndicesBoost();

src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java

@@ -77,6 +77,7 @@ import static org.springframework.data.elasticsearch.utils.IndexBuilder.*;
  * @author Ilkang Na
  * @author Alen Turkovic
  * @author Sascha Woo
+ * @author Jean-Baptiste Nizet
  */
 @RunWith(SpringJUnit4ClassRunner.class)
 @ContextConfiguration("classpath:elasticsearch-template-test.xml")
@@ -1415,6 +1416,88 @@ public class ElasticsearchTemplateTests {
 		});
 	}
 
+	@Test // DATAES-479
+	public void shouldHonorTheHighlightBuilderOptions() {
+
+		// given
+		String documentId = randomNumeric(5);
+		String actualMessage = "some test message with <html> unsafe <script> text";
+		String highlightedMessage = "some <em>test</em> message with &lt;html&gt; unsafe &lt;script&gt; text";
+
+		SampleEntity sampleEntity = SampleEntity.builder()
+				.id(documentId)
+				.message(actualMessage)
+				.version(System.currentTimeMillis())
+				.build();
+
+		IndexQuery indexQuery = getIndexQuery(sampleEntity);
+
+		elasticsearchTemplate.index(indexQuery);
+		elasticsearchTemplate.refresh(SampleEntity.class);
+
+		SearchQuery searchQuery = new NativeSearchQueryBuilder()
+				.withQuery(termQuery("message", "test"))
+				.withHighlightBuilder(new HighlightBuilder().encoder("html"))
+				.withHighlightFields(new HighlightBuilder.Field("message"))
+				.build();
+		// when
+		elasticsearchTemplate.queryForPage(searchQuery, SampleEntity.class, new SearchResultMapper() {
+			@Override
+			public <T> AggregatedPage<T> mapResults(SearchResponse response, Class<T> clazz, Pageable pageable) {
+				for (SearchHit searchHit : response.getHits()) {
+					Map<String, HighlightField> highlightFields = searchHit.getHighlightFields();
+					HighlightField highlightFieldMessage = highlightFields.get("message");
+
+					// then
+					assertNotNull(highlightFieldMessage);
+					assertThat(highlightFieldMessage.fragments()[0].toString(), is(highlightedMessage));
+				}
+				return null;
+			}
+		});
+	}
+
+	@Test // DATAES-479
+	public void shouldHighlightIfBuilderSetEvenIfFieldsNotSet() {
+
+		// given
+		String documentId = randomNumeric(5);
+		String actualMessage = "some test message text";
+		String highlightedMessage = "some <em>test</em> message text";
+
+		SampleEntity sampleEntity = SampleEntity.builder()
+				.id(documentId)
+				.message(actualMessage)
+				.version(System.currentTimeMillis())
+				.build();
+
+		IndexQuery indexQuery = getIndexQuery(sampleEntity);
+
+		elasticsearchTemplate.index(indexQuery);
+		elasticsearchTemplate.refresh(SampleEntity.class);
+
+		SearchQuery searchQuery = new NativeSearchQueryBuilder()
+				.withQuery(termQuery("message", "test"))
+				.withHighlightBuilder(new HighlightBuilder().field("message"))
+				.withHighlightFields(new HighlightBuilder.Field("message"))
+				.build();
+		// when
+		elasticsearchTemplate.queryForPage(searchQuery, SampleEntity.class, new SearchResultMapper() {
+			@Override
+			public <T> AggregatedPage<T> mapResults(SearchResponse response, Class<T> clazz, Pageable pageable) {
+				for (SearchHit searchHit : response.getHits()) {
+					Map<String, HighlightField> highlightFields = searchHit.getHighlightFields();
+					HighlightField highlightFieldMessage = highlightFields.get("message");
+
+					// then
+					assertNotNull(highlightFieldMessage);
+					assertThat(highlightFieldMessage.fragments()[0].toString(), is(highlightedMessage));
+				}
+				return null;
+			}
+		});
+	}
+
 	@Test
 	public void shouldDeleteDocumentBySpecifiedTypeUsingDeleteQuery() {
 		// given
@@ -1538,7 +1621,7 @@ public class ElasticsearchTemplateTests {
 
 	@Test // DATAES-462
 	public void shouldReturnScores() {
-		
+
 		// given
 		List<IndexQuery> indexQueries = new ArrayList<>();
 
@@ -2250,4 +2333,4 @@ public class ElasticsearchTemplateTests {
 			this.lastName = lastName;
 		}
 	}
-}
+}