DATAES-595 - Support for setting preference parameter in query.

Original PR: #288

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

@@ -138,6 +138,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
  * @author Lorenzo Spinelli
  * @author Dmitriy Yakovlev
  * @author Roman Puchkovskiy
+ * @author Farid Azaza
  */
 public class ElasticsearchRestTemplate
 		implements ElasticsearchOperations, EsClient<RestHighLevelClient>, ApplicationContextAware {
@@ -1282,6 +1283,10 @@ public class ElasticsearchRestTemplate
 		if (query.getMinScore() > 0) {
 			sourceBuilder.minScore(query.getMinScore());
 		}
+
+		if (query.getPreference() != null) {
+			request.preference(query.getPreference());
+		}
 		request.source(sourceBuilder);
 		return request;
 	}

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

@@ -121,6 +121,7 @@ import org.springframework.util.StringUtils;
  * @author Christoph Strobl
  * @author Dmitriy Yakovlev
  * @author Peter-Josef Meisch
+ * @author Farid Azaza
  */
 public class ElasticsearchTemplate implements ElasticsearchOperations, EsClient<Client>, ApplicationContextAware {
 
@@ -1069,6 +1070,10 @@ public class ElasticsearchTemplate implements ElasticsearchOperations, EsClient<
 		if (query.getMinScore() > 0) {
 			searchRequestBuilder.setMinScore(query.getMinScore());
 		}
+
+		if (query.getPreference() != null) {
+			searchRequestBuilder.setPreference(query.getPreference());
+		}
 		return searchRequestBuilder;
 	}
 

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

@@ -76,6 +76,7 @@ import org.springframework.util.Assert;
 /**
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Farid Azaza
  * @since 3.2
  */
 public class ReactiveElasticsearchTemplate implements ReactiveElasticsearchOperations {
@@ -273,6 +274,10 @@ public class ReactiveElasticsearchTemplate implements ReactiveElasticsearchOpera
 				request.indicesOptions(query.getIndicesOptions());
 			}
 
+			if (query.getPreference() != null) {
+				request.preference(query.getPreference());
+			}
+
 			Pageable pageable = query.getPageable();
 
 			if (pageable.isPaged()) {

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

@@ -35,6 +35,7 @@ import org.springframework.util.Assert;
  * @author Mark Paluch
  * @author Alen Turkovic
  * @author Sascha Woo
+ * @author Farid Azaza
  */
 abstract class AbstractQuery implements Query {
 
@@ -50,6 +51,7 @@ abstract class AbstractQuery implements Query {
 	protected SearchType searchType = SearchType.DFS_QUERY_THEN_FETCH;
 	protected IndicesOptions indicesOptions;
 	protected boolean trackScores;
+	protected String preference;
 
 	@Override
 	public Sort getSort() {
@@ -63,9 +65,9 @@ abstract class AbstractQuery implements Query {
 
 	@Override
 	public final <T extends Query> T setPageable(Pageable pageable) {
-		
+
 		Assert.notNull(pageable, "Pageable must not be null!");
-		
+
 		this.pageable = pageable;
 		return (T) this.addSort(pageable.getSort());
 	}
@@ -183,4 +185,14 @@ abstract class AbstractQuery implements Query {
 	public void setTrackScores(boolean trackScores) {
 		this.trackScores = trackScores;
 	}
+
+	@Override
+	public String getPreference() {
+		return preference;
+	}
+
+	@Override
+	public void setPreference(String preference) {
+		this.preference = preference;
+	}
 }

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

@@ -38,6 +38,7 @@ import org.springframework.data.elasticsearch.core.facet.FacetRequest;
  * @author Alen Turkovic
  * @author Sascha Woo
  * @author Jean-Baptiste Nizet
+ * @author Farid Azaza
  */
 public class NativeSearchQueryBuilder {
 
@@ -61,6 +62,7 @@ public class NativeSearchQueryBuilder {
 	private String route;
 	private SearchType searchType;
 	private IndicesOptions indicesOptions;
+	private String preference;
 
 	public NativeSearchQueryBuilder withQuery(QueryBuilder queryBuilder) {
 		this.queryBuilder = queryBuilder;
@@ -167,6 +169,11 @@ public class NativeSearchQueryBuilder {
 		return this;
 	}
 
+	public NativeSearchQueryBuilder withPreference(String preference) {
+		this.preference = preference;
+		return this;
+	}
+
 	public NativeSearchQuery build() {
 		NativeSearchQuery nativeSearchQuery = new NativeSearchQuery(queryBuilder, filterBuilder, sortBuilders,
 				highlightBuilder, highlightFields);
@@ -225,7 +232,9 @@ public class NativeSearchQueryBuilder {
 		if (indicesOptions != null) {
 			nativeSearchQuery.setIndicesOptions(indicesOptions);
 		}
-
+		if (preference != null) {
+			nativeSearchQuery.setPreference(preference);
+		}
 		return nativeSearchQuery;
 	}
 }

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

@@ -35,6 +35,7 @@ import org.springframework.data.domain.Sort;
  * @author Alen Turkovic
  * @author Sascha Woo
  * @author Christoph Strobl
+ * @author Farid Azaza
  */
 public interface Query {
 
@@ -178,4 +179,18 @@ public interface Query {
 	 * @return null if not set
 	 */
 	IndicesOptions getIndicesOptions();
+
+	/**
+	 * Get preference
+	 *
+	 * @return
+	 */
+	String getPreference();
+
+	/**
+	 * Add preference filter to be added as part of search request
+	 *
+	 * @param preference
+	 */
+	void setPreference(String preference);
 }

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

@@ -106,6 +106,7 @@ import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
  * @author Ivan Greene
  * @author Dmitriy Yakovlev
  * @author Peter-Josef Meisch
+ * @author Farid Azaza
  */
 @RunWith(SpringJUnit4ClassRunner.class)
 @ContextConfiguration("classpath:elasticsearch-template-test.xml")
@@ -330,6 +331,54 @@ public class ElasticsearchTemplateTests {
 		assertThat(sampleEntities.getTotalElements()).isGreaterThanOrEqualTo(1);
 	}
 
+	@Test // DATAES-595
+	public void shouldReturnPageUsingLocalPreferenceForGivenSearchQuery() {
+
+		// given
+		String documentId = randomNumeric(5);
+		SampleEntity sampleEntity = SampleEntity.builder().id(documentId).message("some message")
+				.version(System.currentTimeMillis()).build();
+
+		IndexQuery indexQuery = getIndexQuery(sampleEntity);
+
+		elasticsearchTemplate.index(indexQuery);
+		elasticsearchTemplate.refresh(SampleEntity.class);
+
+		SearchQuery searchQueryWithValidPreference = new NativeSearchQueryBuilder().withQuery(matchAllQuery())
+				.withPreference("_local").build();
+
+		// when
+		Page<SampleEntity> sampleEntities = elasticsearchTemplate.queryForPage(searchQueryWithValidPreference,
+				SampleEntity.class);
+
+		// then
+		assertThat(sampleEntities).isNotNull();
+		assertThat(sampleEntities.getTotalElements()).isGreaterThanOrEqualTo(1);
+	}
+
+	@Test(expected = Exception.class) // DATAES-595
+	public void shouldThrowExceptionWhenInvalidPreferenceForSearchQuery() {
+
+		// given
+		String documentId = randomNumeric(5);
+		SampleEntity sampleEntity = SampleEntity.builder().id(documentId).message("some message")
+				.version(System.currentTimeMillis()).build();
+
+		IndexQuery indexQuery = getIndexQuery(sampleEntity);
+
+		elasticsearchTemplate.index(indexQuery);
+		elasticsearchTemplate.refresh(SampleEntity.class);
+
+		SearchQuery searchQueryWithInvalidPreference = new NativeSearchQueryBuilder().withQuery(matchAllQuery())
+				.withPreference("_only_nodes:oops").build();
+
+		// when
+		elasticsearchTemplate.queryForPage(searchQueryWithInvalidPreference, SampleEntity.class);
+
+		// then Throw IllegalArgumentException in case of ElasticsearchTemplate and ElasticsearchStatusException in case of
+		// ElasticsearchRestTemplate
+	}
+
 	@Test // DATAES-422 - Add support for IndicesOptions in search queries
 	public void shouldPassIndicesOptionsForGivenSearchQuery() {
 

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

@@ -39,6 +39,7 @@ import java.util.UUID;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import org.elasticsearch.ElasticsearchStatusException;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -75,6 +76,7 @@ import org.springframework.util.StringUtils;
  * @author Christoph Strobl
  * @author Mark Paluch
  * @author Peter-Josef Meisch
+ * @author Farid Azaza
  * @currentRead Golden Fool - Robin Hobb
  */
 @RunWith(SpringRunner.class)
@@ -389,6 +391,43 @@ public class ReactiveElasticsearchTemplateTests {
 				.verifyComplete();
 	}
 
+	@Test // DATAES-595
+	public void shouldReturnListUsingLocalPreferenceForGivenCriteria() {
+
+		SampleEntity sampleEntity1 = randomEntity("test message");
+		SampleEntity sampleEntity2 = randomEntity("test test");
+		SampleEntity sampleEntity3 = randomEntity("some message");
+
+		index(sampleEntity1, sampleEntity2, sampleEntity3);
+
+		CriteriaQuery queryWithValidPreference = new CriteriaQuery(
+				new Criteria("message").contains("some").and("message").contains("message"));
+		queryWithValidPreference.setPreference("_local");
+
+		template.find(queryWithValidPreference, SampleEntity.class) //
+				.as(StepVerifier::create) //
+				.expectNext(sampleEntity3) //
+				.verifyComplete();
+	}
+
+	@Test // DATAES-595
+	public void shouldThrowElasticsearchStatusExceptionWhenInvalidPreferenceForGivenCriteria() {
+
+		SampleEntity sampleEntity1 = randomEntity("test message");
+		SampleEntity sampleEntity2 = randomEntity("test test");
+		SampleEntity sampleEntity3 = randomEntity("some message");
+
+		index(sampleEntity1, sampleEntity2, sampleEntity3);
+
+		CriteriaQuery queryWithInvalidPreference = new CriteriaQuery(
+				new Criteria("message").contains("some").and("message").contains("message"));
+		queryWithInvalidPreference.setPreference("_only_nodes:oops");
+
+		template.find(queryWithInvalidPreference, SampleEntity.class) //
+				.as(StepVerifier::create) //
+				.expectError(ElasticsearchStatusException.class).verify();
+	}
+
 	@Test // DATAES-504
 	public void shouldReturnProjectedTargetEntity() {