DATAES-42 : Adding support for letting Elasticsearch generate Id for document

mohsinh · mohsinh · commit fb95bb06bf99 · 2014-02-01T12:53:07.000Z
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/DefaultResultMapper.java b/src/main/java/org/springframework/data/elasticsearch/core/DefaultResultMapper.java
@@ -26,9 +26,15 @@
 import org.elasticsearch.search.SearchHitField;
 import org.elasticsearch.search.facet.Facet;
 import org.springframework.data.domain.Pageable;
+import org.springframework.data.elasticsearch.annotations.Document;
 import org.springframework.data.elasticsearch.core.facet.DefaultFacetMapper;
 import org.springframework.data.elasticsearch.core.facet.FacetResult;
+import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
+import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty;
+import org.springframework.data.mapping.PersistentProperty;
+import org.springframework.data.mapping.context.MappingContext;
 
+import java.lang.reflect.Method;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.charset.Charset;
@@ -41,9 +47,16 @@
  */
 public class DefaultResultMapper extends AbstractResultMapper {
 
+    private MappingContext<? extends ElasticsearchPersistentEntity<?>, ElasticsearchPersistentProperty> mappingContext;
+   
     public DefaultResultMapper(){
         super(new DefaultEntityMapper());
     }
+    
+    public DefaultResultMapper(MappingContext<? extends ElasticsearchPersistentEntity<?>, ElasticsearchPersistentProperty> mappingContext){
+       super(new DefaultEntityMapper());
+       this.mappingContext = mappingContext;
+    }
 
     public DefaultResultMapper(EntityMapper entityMapper) {
         super(entityMapper);
@@ -55,11 +68,14 @@ public <T> FacetedPage<T> mapResults(SearchResponse response, Class<T> clazz, Pa
         List<T> results = new ArrayList<T>();
         for (SearchHit hit : response.getHits()) {
             if (hit != null) {
+                T result = null;
                 if (!Strings.isNullOrEmpty(hit.sourceAsString())) {
-                    results.add(mapEntity(hit.sourceAsString(), clazz));
+                    result = mapEntity(hit.sourceAsString(), clazz);
                 } else {
-                    results.add(mapEntity(hit.getFields().values(), clazz));
+                    result = mapEntity(hit.getFields().values(), clazz);
                 }
+                setPersistentEntityId(result, hit.getId(), clazz);
+                results.add(result);
             }
         }
         List<FacetResult> facets = new ArrayList<FacetResult>();
@@ -106,6 +122,27 @@ private String buildJSONFromFields(Collection<SearchHitField> values) {
 
     @Override
     public <T> T mapResult(GetResponse response, Class<T> clazz) {
-        return mapEntity(response.getSourceAsString(),clazz);
+        T result = mapEntity(response.getSourceAsString(),clazz);
+        if (result != null){
+           setPersistentEntityId(result, response.getId(), clazz);
+        }
+        return result;
+    }
+    
+    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();
+          // Only deal with String because ES generated Ids are strings !
+          if (idProperty != null && idProperty.getType().isAssignableFrom(String.class)){
+              Method setter = idProperty.getSetter();
+              if (setter != null){
+                  try{
+                      setter.invoke(result, id);
+                  } catch (Throwable t) {
+                      t.printStackTrace();
+                  }
+              }
+          }
+      }
     }
 }
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java
@@ -56,9 +56,11 @@
 import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
 import org.springframework.data.elasticsearch.core.mapping.SimpleElasticsearchMappingContext;
 import org.springframework.data.elasticsearch.core.query.*;
+import org.springframework.data.mapping.PersistentProperty;
 import org.springframework.util.Assert;
 
 import java.io.IOException;
+import java.lang.reflect.Method;
 import java.util.*;
 
 import static org.apache.commons.collections.CollectionUtils.isNotEmpty;
@@ -106,7 +108,7 @@ public ElasticsearchTemplate(Client client, ElasticsearchConverter elasticsearch
         this.client = client;
         this.elasticsearchConverter = (elasticsearchConverter == null) ? new MappingElasticsearchConverter(
                 new SimpleElasticsearchMappingContext()) : elasticsearchConverter;
-        this.resultsMapper = (resultsMapper == null) ? new DefaultResultMapper() : resultsMapper;
+        this.resultsMapper = (resultsMapper == null) ? new DefaultResultMapper(this.elasticsearchConverter.getMappingContext()) : resultsMapper;
     }
 
     @Override
@@ -145,7 +147,9 @@ public <T> T queryForObject(GetQuery query, Class<T> clazz, GetResultMapper mapp
         GetResponse response = client
                 .prepareGet(persistentEntity.getIndexName(), persistentEntity.getIndexType(), query.getId()).execute()
                 .actionGet();
-        return mapper.mapResult(response, clazz);
+
+        T entity = mapper.mapResult(response, clazz);
+        return entity;
     }
 
     @Override
@@ -246,7 +250,12 @@ public <T> long count(SearchQuery query, Class<T> clazz) {
 
     @Override
     public String index(IndexQuery query) {
-        return prepareIndex(query).execute().actionGet().getId();
+        String documentId = prepareIndex(query).execute().actionGet().getId();
+        // We should call this because we are not going through a mapper.
+        if (query.getObject() != null){
+           setPersistentEntityId(query.getObject(), documentId);
+        }
+        return documentId;
     }
 
     @Override
@@ -540,10 +549,16 @@ private IndexRequestBuilder prepareIndex(IndexQuery query) {
 
             IndexRequestBuilder indexRequestBuilder = null;
 
-            if(query.getObject() != null) {
-                indexRequestBuilder = client.prepareIndex(indexName, type, query.getId()).setSource(
-                    resultsMapper.getEntityMapper().mapToString(query.getObject()));
-            } else if(query.getSource() != null) {
+            if (query.getObject() != null) {
+                // If we have a query id and a document id, do not ask ES to generate one.
+                String entityId = getPersistentEntityId(query.getObject());
+                if (query.getId() != null && entityId != null){
+                   indexRequestBuilder = client.prepareIndex(indexName, type, query.getId());   
+                } else {
+                   indexRequestBuilder = client.prepareIndex(indexName, type);   
+                }
+                indexRequestBuilder.setSource(resultsMapper.getEntityMapper().mapToString(query.getObject()));
+            } else if (query.getSource() != null) {
                 indexRequestBuilder = client.prepareIndex(indexName, type, query.getId()).setSource(query.getSource());
             } else {
                 throw new ElasticsearchException("object or source is null, failed to index the document [id: " + query.getId() + "]");
@@ -613,6 +628,40 @@ private ElasticsearchPersistentEntity getPersistentEntityFor(Class clazz) {
                 + " is not a Document. Make sure the document class is annotated with @Document(indexName=\"foo\")");
         return elasticsearchConverter.getMappingContext().getPersistentEntity(clazz);
     }
+    
+    private String getPersistentEntityId(Object entity){
+       PersistentProperty idProperty = getPersistentEntityFor(entity.getClass()).getIdProperty();
+       if (idProperty != null){
+          Method getter = idProperty.getGetter();
+          if (getter != null){
+             try{
+                Object id = getter.invoke(entity);
+                if (id != null){
+                   return String.valueOf(id);
+                }
+                
+             } catch (Throwable t){
+                t.printStackTrace();
+             }
+          }
+       }
+       return null;
+    }
+    
+    private void setPersistentEntityId(Object entity, String id){
+       PersistentProperty idProperty = getPersistentEntityFor(entity.getClass()).getIdProperty();
+       // Only deal with String because ES generated Ids are strings !
+       if (idProperty != null && idProperty.getType().isAssignableFrom(String.class)){
+          Method setter = idProperty.getSetter();
+          if (setter != null){
+             try{
+                setter.invoke(entity, id);
+             } catch (Throwable t) {
+                t.printStackTrace();
+             }
+          }
+       }
+    }
 
     private String[] retrieveIndexNameFromPersistentEntity(Class clazz) {
         return new String[]{getPersistentEntityFor(clazz).getIndexName()};
diff --git a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java b/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java
@@ -1034,4 +1034,62 @@ public void shouldReturnDocumentAboveMinimalScoreGivenQuery() {
         assertThat(page.getTotalElements(),is(1L));
         assertThat(page.getContent().get(0).getMessage(), is("ab"));
     }
+
+
+    @Test
+    public void shouldDoIndexWithoutId() {
+        // given
+        // document
+        SampleEntity sampleEntity = new SampleEntity();
+        sampleEntity.setMessage("some message");
+        sampleEntity.setVersion(System.currentTimeMillis());
+
+        IndexQuery indexQuery = new IndexQuery();
+        indexQuery.setObject(sampleEntity);
+        // when
+        String documentId = elasticsearchTemplate.index(indexQuery);
+        // then
+        assertThat(sampleEntity.getId(), is(equalTo(documentId)));
+
+        GetQuery getQuery = new GetQuery();
+        getQuery.setId(documentId);
+        SampleEntity result = elasticsearchTemplate.queryForObject(getQuery, SampleEntity.class);
+        assertThat(result.getId(), is(equalTo(documentId)));
+    }
+
+    @Test
+    public void shouldDoBulkIndexWithoutId() {
+        // given
+        List<IndexQuery> indexQueries = new ArrayList<IndexQuery>();
+        // first document
+        String documentId = randomNumeric(5);
+        SampleEntity sampleEntity1 = new SampleEntity();
+        sampleEntity1.setMessage("some message");
+        sampleEntity1.setVersion(System.currentTimeMillis());
+
+        IndexQuery indexQuery1 = new IndexQuery();
+        //indexQuery1.setId(documentId);
+        indexQuery1.setObject(sampleEntity1);
+        indexQueries.add(indexQuery1);
+
+        // second document
+        SampleEntity sampleEntity2 = new SampleEntity();
+        sampleEntity2.setMessage("some message");
+        sampleEntity2.setVersion(System.currentTimeMillis());
+
+        IndexQuery indexQuery2 = new IndexQuery();
+        indexQuery2.setObject(sampleEntity2);
+        indexQueries.add(indexQuery2);
+        // when
+        elasticsearchTemplate.bulkIndex(indexQueries);
+        elasticsearchTemplate.refresh(SampleEntity.class, true);
+        // then
+        SearchQuery searchQuery = new NativeSearchQueryBuilder().withQuery(matchAllQuery()).build();
+        Page<SampleEntity> sampleEntities = elasticsearchTemplate.queryForPage(searchQuery, SampleEntity.class);
+        assertThat(sampleEntities.getTotalElements(), is(equalTo(2L)));
+
+        assertThat(sampleEntities.getContent().get(0).getId(), is(notNullValue()));
+        assertThat(sampleEntities.getContent().get(1).getId(), is(notNullValue()));
+    }
+
 }
diff --git a/src/test/java/org/springframework/data/elasticsearch/repository/support/IntegerIDRepositoryTests.java b/src/test/java/org/springframework/data/elasticsearch/repository/support/IntegerIDRepositoryTests.java
@@ -20,7 +20,6 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.data.elasticsearch.DoubleIDEntity;
 import org.springframework.data.elasticsearch.IntegerIDEntity;
 import org.springframework.data.elasticsearch.core.ElasticsearchTemplate;
 import org.springframework.data.elasticsearch.repositories.IntegerIDRepository;
diff --git a/src/test/java/org/springframework/data/elasticsearch/repository/support/SimpleElasticsearchRepositoryTests.java b/src/test/java/org/springframework/data/elasticsearch/repository/support/SimpleElasticsearchRepositoryTests.java
@@ -104,6 +104,18 @@ public void shouldSaveDocument() {
 		assertThat(entityFromElasticSearch, is(notNullValue()));
 	}
 
+	@Test
+   public void shouldSaveDocumentWithoutId() {
+      // given
+      SampleEntity sampleEntity = new SampleEntity();
+      sampleEntity.setMessage("some message");
+      sampleEntity.setVersion(System.currentTimeMillis());
+      // when
+      repository.save(sampleEntity);
+      // then
+      assertThat(sampleEntity.getId(), is(notNullValue()));
+   }
+
 	@Test
 	public void shouldFindDocumentById() {
 		// given
