DATAES-188 - Source filtering feature Implementation

Author: mohsinh

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

@@ -954,6 +954,11 @@ private SearchRequestBuilder prepareSearch(Query query) {
 		SearchRequestBuilder searchRequestBuilder = client.prepareSearch(toArray(query.getIndices()))
 				.setSearchType(query.getSearchType()).setTypes(toArray(query.getTypes()));
 
+		if (query.getSourceFilter() != null) {
+			SourceFilter sourceFilter = query.getSourceFilter();
+			searchRequestBuilder.setFetchSource(sourceFilter.getIncludes(), sourceFilter.getExcludes());
+		}
+
 		if (query.getPageable() != null) {
 			startRecord = query.getPageable().getPageNumber() * query.getPageable().getPageSize();
 			searchRequestBuilder.setSize(query.getPageable().getPageSize());

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

@@ -39,6 +39,7 @@ abstract class AbstractQuery implements Query {
 	protected List<String> indices = new ArrayList<String>();
 	protected List<String> types = new ArrayList<String>();
 	protected List<String> fields = new ArrayList<String>();
+	protected SourceFilter sourceFilter;
 	protected float minScore;
 	protected Collection<String> ids;
 	protected String route;
@@ -91,6 +92,16 @@ public List<String> getTypes() {
 		return types;
 	}
 
+	@Override
+	public void addSourceFilter(SourceFilter sourceFilter) {
+		this.sourceFilter = sourceFilter;
+	}
+
+	@Override
+	public SourceFilter getSourceFilter() {
+		return sourceFilter;
+	}
+
 	@SuppressWarnings("unchecked")
 	public final <T extends Query> T addSort(Sort sort) {
 		if (sort == null) {

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2016 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.elasticsearch.core.query;
+
+/**
+ * SourceFilter implementation for providing includes and excludes.
+ *
+ * @Author Jon Tsiros
+ */
+public class FetchSourceFilter implements SourceFilter {
+
+	private final String[] includes;
+	private final String[] excludes;
+
+	public FetchSourceFilter(final String[] includes, final String[] excludes) {
+		this.includes = includes;
+		this.excludes = excludes;
+	}
+
+	@Override
+	public String[] getIncludes() {
+		return includes;
+	}
+
+	@Override
+	public String[] getExcludes() {
+		return excludes;
+	}
+}

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

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2016 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.elasticsearch.core.query;
+
+/**
+ * SourceFilter builder for providing includes and excludes.
+ *
+ * @Author Jon Tsiros
+ */
+public class FetchSourceFilterBuilder {
+
+	private String[] includes;
+	private String[] excludes;
+
+	public FetchSourceFilterBuilder withIncludes(String... includes) {
+		this.includes = includes;
+		return this;
+	}
+
+	public FetchSourceFilterBuilder withExcludes(String... excludes) {
+		this.excludes = excludes;
+		return this;
+	}
+
+	public SourceFilter build() {
+		if (includes == null) includes = new String[0];
+		if (excludes == null) excludes = new String[0];
+
+		SourceFilter sourceFilter = new FetchSourceFilter(includes, excludes);
+		return sourceFilter;
+	}
+}

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

@@ -41,13 +41,13 @@ public class NativeSearchQueryBuilder {
 	private QueryBuilder filterBuilder;
     private List<ScriptField> scriptFields = new ArrayList<ScriptField>();
 	private List<SortBuilder> sortBuilders = new ArrayList<SortBuilder>();
-	/*private List<FacetRequest> facetRequests = new ArrayList<FacetRequest>();*/
 	private List<AbstractAggregationBuilder> aggregationBuilders = new ArrayList<AbstractAggregationBuilder>();
 	private HighlightBuilder.Field[] highlightFields;
 	private Pageable pageable;
 	private String[] indices;
 	private String[] types;
 	private String[] fields;
+	private SourceFilter sourceFilter;
 	private List<IndexBoost> indicesBoost;
 	private float minScore;
 	private Collection<String> ids;
@@ -114,6 +114,11 @@ public NativeSearchQueryBuilder withFields(String... fields) {
 		return this;
 	}
 
+	public NativeSearchQueryBuilder withSourceFilter(SourceFilter sourceFilter) {
+				this.sourceFilter = sourceFilter;
+				return this;
+	}
+
 	public NativeSearchQueryBuilder withMinScore(float minScore) {
 		this.minScore = minScore;
 		return this;
@@ -151,6 +156,10 @@ public NativeSearchQuery build() {
 		if (fields != null) {
 			nativeSearchQuery.addFields(fields);
 		}
+
+		if (sourceFilter != null) {
+			nativeSearchQuery.addSourceFilter(sourceFilter);
+		}
 
 		if(indicesBoost != null) {
 		    nativeSearchQuery.setIndicesBoost(indicesBoost);
@@ -160,10 +169,6 @@ public NativeSearchQuery build() {
             nativeSearchQuery.setScriptFields(scriptFields);
         }
 
-/*		if (CollectionUtils.isNotEmpty(facetRequests)) {
-			nativeSearchQuery.setFacets(facetRequests);
-		}*/
-
 		if (CollectionUtils.isNotEmpty(aggregationBuilders)) {
 			nativeSearchQuery.setAggregations(aggregationBuilders);
 		}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2016 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.
@@ -111,6 +111,21 @@ public interface Query {
 	 */
 	List<String> getFields();
 
+	/**
+	 * Add source filter to be added as part of search request
+	 *
+	 * @param sourceFilter
+	 */
+	void addSourceFilter(SourceFilter sourceFilter);
+
+	/**
+	 * Get SourceFilter to be returned to get include and exclude source
+	 * fields as part of search request.
+	 *
+	 * @return SourceFilter
+	 */
+	SourceFilter getSourceFilter();
+
 	/**
 	 * Get minimum score
 	 *

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

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2016 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.elasticsearch.core.query;
+
+/**
+ * SourceFilter for providing includes and excludes.
+ *
+ * @Author Jon Tsiros
+ */
+public interface SourceFilter {
+
+	String[] getIncludes();
+
+	String[] getExcludes();
+}

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

@@ -489,8 +489,8 @@ public void shouldUseScriptedFields() {
 		SearchQuery searchQuery = new NativeSearchQueryBuilder()
 				.withQuery(matchAllQuery())
 				.withScriptField(new ScriptField("scriptedRate",
-						new Script("doc['rate'].value * factor", ScriptService.ScriptType.INLINE, null , params)))
-						.build();
+						new Script("doc['rate'].value * factor", ScriptService.ScriptType.INLINE, null, params)))
+				.build();
 		Page<SampleEntity> sampleEntities = elasticsearchTemplate.queryForPage(searchQuery, SampleEntity.class);
 		// then
 		assertThat(sampleEntities.getTotalElements(), equalTo(1L));
@@ -646,6 +646,33 @@ public <T> Page<T> mapResults(SearchResponse response, Class<T> clazz, Pageable
 		assertThat(page.getContent().get(0), is(message));
 	}
 
+	@Test
+	public void shouldReturnFieldsBasedOnSourceFilter() {
+		// given
+		String documentId = randomNumeric(5);
+		String message = "some test message";
+		SampleEntity sampleEntity = SampleEntity.builder().id(documentId).message(message)
+				.version(System.currentTimeMillis()).build();
+
+		IndexQuery indexQuery = getIndexQuery(sampleEntity);
+
+		elasticsearchTemplate.index(indexQuery);
+		elasticsearchTemplate.refresh(SampleEntity.class);
+
+		FetchSourceFilterBuilder sourceFilter = new FetchSourceFilterBuilder();
+		sourceFilter.withIncludes("message");
+
+		SearchQuery searchQuery = new NativeSearchQueryBuilder().withQuery(matchAllQuery()).withIndices(INDEX_NAME)
+				.withTypes(TYPE_NAME).withSourceFilter(sourceFilter.build()).build();
+		// when
+		Page<SampleEntity> page = elasticsearchTemplate.queryForPage(searchQuery, SampleEntity.class);
+		// then
+		assertThat(page, is(notNullValue()));
+		assertThat(page.getTotalElements(), is(equalTo(1L)));
+		assertThat(page.getContent().get(0).getMessage(), is(message));
+	}
+
+
 	@Test
 	public void shouldReturnSimilarResultsGivenMoreLikeThisQuery() {
 		// given