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

This commit is contained in:
mohsin.husen 2014-02-01 12:53:07 +00:00
parent bbc46d95df
commit fb95bb06bf
5 changed files with 166 additions and 11 deletions

View File

@ -26,9 +26,15 @@ import org.elasticsearch.search.SearchHit;
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 @@ import java.util.List;
*/
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 class DefaultResultMapper extends AbstractResultMapper {
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 @@ public class DefaultResultMapper extends AbstractResultMapper {
@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();
}
}
}
}
}
}

View File

@ -56,9 +56,11 @@ import org.springframework.data.elasticsearch.core.facet.FacetRequest;
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 class ElasticsearchTemplate implements ElasticsearchOperations {
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 class ElasticsearchTemplate implements ElasticsearchOperations {
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 class ElasticsearchTemplate implements ElasticsearchOperations {
@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 @@ public class ElasticsearchTemplate implements ElasticsearchOperations {
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 @@ public class ElasticsearchTemplate implements ElasticsearchOperations {
+ " 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()};

View File

@ -1034,4 +1034,62 @@ public class ElasticsearchTemplateTests {
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()));
}
}

View File

@ -20,7 +20,6 @@ import org.junit.Before;
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;

View File

@ -104,6 +104,18 @@ public class SimpleElasticsearchRepositoryTests {
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