DATAES-52 - Support Get & Multi Get

Added support for MultiGetResultMapper for custom fields

Author: mohsinh

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

@@ -22,9 +22,12 @@
 import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.LinkedList;
 import java.util.List;
 
 import org.elasticsearch.action.get.GetResponse;
+import org.elasticsearch.action.get.MultiGetItemResponse;
+import org.elasticsearch.action.get.MultiGetResponse;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.common.base.Strings;
 import org.elasticsearch.common.jackson.core.JsonEncoding;
@@ -129,6 +132,18 @@ public <T> T mapResult(GetResponse response, Class<T> clazz) {
 		return result;
 	}
 
+	@Override
+	public <T> LinkedList<T> mapResults(MultiGetResponse responses, Class<T> clazz) {
+		LinkedList<T> list = new LinkedList<T>();
+		for (MultiGetItemResponse response : responses.getResponses()) {
+			if (!response.isFailed() && response.getResponse().isExists()) {
+				T result = mapEntity(response.getResponse().getSourceAsString(), clazz);
+				list.add(result);
+			}
+		}
+		return list;
+	}
+
 	private <T> void setPersistentEntityId(T result, String id, Class<T> clazz) {
 		if (mappingContext != null && clazz.isAnnotationPresent(Document.class)) {
 			PersistentProperty<ElasticsearchPersistentProperty> idProperty = mappingContext.getPersistentEntity(clazz).getIdProperty();

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

@@ -183,14 +183,24 @@ public interface ElasticsearchOperations {
 	<T> long count(SearchQuery query, Class<T> clazz);
 
 	/**
-	 * Execute a multiget against elasticsearch for the given ids
+	 * Execute a multiGet against elasticsearch for the given ids
 	 *
 	 * @param searchQuery
 	 * @param clazz
 	 * @return
 	 */
 	<T> LinkedList<T> multiGet(SearchQuery searchQuery, Class<T> clazz);
 
+	/**
+	 * Execute a multiGet against elasticsearch for the given ids with MultiGetResultMapper
+	 *
+	 * @param searchQuery
+	 * @param clazz
+	 * @param multiGetResultMapper
+	 * @return
+	 */
+	<T> LinkedList<T> multiGet(SearchQuery searchQuery, Class<T> clazz, MultiGetResultMapper multiGetResultMapper);
+
 	/**
 	 * Index an object. Will do save or update
 	 *

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

@@ -37,7 +37,10 @@
 import org.elasticsearch.action.bulk.BulkRequestBuilder;
 import org.elasticsearch.action.bulk.BulkResponse;
 import org.elasticsearch.action.count.CountRequestBuilder;
-import org.elasticsearch.action.get.*;
+import org.elasticsearch.action.get.GetResponse;
+import org.elasticsearch.action.get.MultiGetRequest;
+import org.elasticsearch.action.get.MultiGetRequestBuilder;
+import org.elasticsearch.action.get.MultiGetResponse;
 import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.action.mlt.MoreLikeThisRequestBuilder;
 import org.elasticsearch.action.search.SearchRequestBuilder;
@@ -248,6 +251,10 @@ public <T> long count(SearchQuery query, Class<T> clazz) {
 
 	@Override
 	public <T> LinkedList<T> multiGet(SearchQuery searchQuery, Class<T> clazz) {
+		return resultsMapper.mapResults(getMultiResponse(searchQuery, clazz), clazz);
+	}
+
+	private <T> MultiGetResponse getMultiResponse(SearchQuery searchQuery, Class<T> clazz) {
 
 		ElasticsearchPersistentEntity<T> persistentEntity = getPersistentEntityFor(clazz);
 
@@ -266,14 +273,12 @@ public <T> LinkedList<T> multiGet(SearchQuery searchQuery, Class<T> clazz) {
 			}
 			builder.add(item);
 		}
-		MultiGetResponse responses = builder.execute().actionGet();
-		final LinkedList<T> result = new LinkedList<T>();
-		for (MultiGetItemResponse response : responses.getResponses()) {
-			if (!response.isFailed() && response.getResponse().isExists()) {
-				result.add(resultsMapper.mapResult(response.getResponse(), clazz));
-			}
-		}
-		return result;
+		return builder.execute().actionGet();
+	}
+
+	@Override
+	public <T> LinkedList<T> multiGet(SearchQuery searchQuery, Class<T> clazz, MultiGetResultMapper getResultMapper) {
+		return getResultMapper.mapResults(getMultiResponse(searchQuery, clazz), clazz);
 	}
 
 	@Override

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

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2014 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;
+
+import java.util.LinkedList;
+
+import org.elasticsearch.action.get.MultiGetResponse;
+
+/**
+ * @author Mohsin Husen
+ */
+public interface MultiGetResultMapper {
+
+	<T> LinkedList<T> mapResults(MultiGetResponse responses, Class<T> clazz);
+}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2014 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.
@@ -23,7 +23,7 @@
  * @author Artur Konczak
  */
 
-public interface ResultsMapper extends SearchResultMapper, GetResultMapper {
+public interface ResultsMapper extends SearchResultMapper, GetResultMapper, MultiGetResultMapper {
 
 	EntityMapper getEntityMapper();
 }

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

@@ -15,7 +15,10 @@
  */
 package org.springframework.data.elasticsearch.core;
 
+import java.util.LinkedList;
+
 import org.elasticsearch.action.get.GetResponse;
+import org.elasticsearch.action.get.MultiGetResponse;
 import org.elasticsearch.action.search.SearchResponse;
 import org.springframework.data.domain.Pageable;
 
@@ -46,4 +49,9 @@ public <T> T mapResult(GetResponse response, Class<T> clazz) {
 	public <T> FacetedPage<T> mapResults(SearchResponse response, Class<T> clazz, Pageable pageable) {
 		return null;  //To change body of implemented methods use File | Settings | File Templates.
 	}
+
+	@Override
+	public <T> LinkedList<T> mapResults(MultiGetResponse responses, Class<T> clazz) {
+		return null;
+	}
 }

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

@@ -23,6 +23,8 @@
 
 import java.util.*;
 
+import org.elasticsearch.action.get.MultiGetItemResponse;
+import org.elasticsearch.action.get.MultiGetResponse;
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.index.engine.DocumentMissingException;
@@ -153,6 +155,63 @@ public void shouldReturnObjectsForGivenIdsUsingMultiGet() {
 		assertEquals(sampleEntities.get(1), sampleEntity2);
 	}
 
+	@Test
+	public void shouldReturnObjectsForGivenIdsUsingMultiGetWithFields() {
+		// given
+		List<IndexQuery> indexQueries = new ArrayList<IndexQuery>();
+		// first document
+		String documentId = randomNumeric(5);
+		SampleEntity sampleEntity1 = new SampleEntity();
+		sampleEntity1.setId(documentId);
+		sampleEntity1.setMessage("some message");
+		sampleEntity1.setType("type1");
+		sampleEntity1.setVersion(System.currentTimeMillis());
+
+		IndexQuery indexQuery1 = new IndexQuery();
+		indexQuery1.setId(documentId);
+		indexQuery1.setObject(sampleEntity1);
+		indexQueries.add(indexQuery1);
+
+		// second document
+		String documentId2 = randomNumeric(5);
+		SampleEntity sampleEntity2 = new SampleEntity();
+		sampleEntity2.setId(documentId2);
+		sampleEntity2.setMessage("some message");
+		sampleEntity2.setType("type2");
+		sampleEntity2.setVersion(System.currentTimeMillis());
+
+		IndexQuery indexQuery2 = new IndexQuery();
+		indexQuery2.setId(documentId2);
+		indexQuery2.setObject(sampleEntity2);
+
+		indexQueries.add(indexQuery2);
+
+		elasticsearchTemplate.bulkIndex(indexQueries);
+		elasticsearchTemplate.refresh(SampleEntity.class, true);
+
+		// when
+		SearchQuery query = new NativeSearchQueryBuilder()
+				.withIds(Arrays.asList(documentId, documentId2))
+				.withFields("message", "type")
+				.build();
+		LinkedList<SampleEntity> sampleEntities = elasticsearchTemplate.multiGet(query, SampleEntity.class, new MultiGetResultMapper() {
+			@Override
+			public <T> LinkedList<T> mapResults(MultiGetResponse responses, Class<T> clazz) {
+				LinkedList<T> list = new LinkedList<T>();
+				for (MultiGetItemResponse response : responses.getResponses()) {
+					SampleEntity entity = new SampleEntity();
+					entity.setId(response.getResponse().getId());
+					entity.setMessage((String) response.getResponse().getField("message").getValue());
+					entity.setType((String) response.getResponse().getField("type").getValue());
+					list.add((T) entity);
+				}
+				return list;
+			}
+		});
+		// then
+		assertThat(sampleEntities.size(), is(equalTo(2)));
+	}
+
 	@Test
 	public void shouldReturnPageForGivenSearchQuery() {
 		// given