diff --git a/src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchPartQuery.java b/src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchPartQuery.java index a0edff13c..0e90250da 100644 --- a/src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchPartQuery.java +++ b/src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchPartQuery.java @@ -37,7 +37,9 @@ import org.springframework.util.ClassUtils; * @author Rasmus Faber-Espensen */ public class ElasticsearchPartQuery extends AbstractElasticsearchRepositoryQuery { + private static final int DEFAULT_STREAM_BATCH_SIZE = 500; + private final PartTree tree; private final MappingContext mappingContext; @@ -49,37 +51,49 @@ public class ElasticsearchPartQuery extends AbstractElasticsearchRepositoryQuery @Override public Object execute(Object[] parameters) { + ParametersParameterAccessor accessor = new ParametersParameterAccessor(queryMethod.getParameters(), parameters); CriteriaQuery query = createQuery(accessor); if (tree.isDelete()) { + Object result = countOrGetDocumentsForDelete(query, accessor); elasticsearchOperations.delete(query, queryMethod.getEntityInformation().getJavaType()); return result; } else if (queryMethod.isPageQuery()) { + query.setPageable(accessor.getPageable()); return elasticsearchOperations.queryForPage(query, queryMethod.getEntityInformation().getJavaType()); } else if (queryMethod.isStreamQuery()) { + Class entityType = queryMethod.getEntityInformation().getJavaType(); if (accessor.getPageable().isUnpaged()) { + query.setPageable(PageRequest.of(0, DEFAULT_STREAM_BATCH_SIZE)); } else { + query.setPageable(accessor.getPageable()); } return StreamUtils .createStreamFromIterator((CloseableIterator) elasticsearchOperations.stream(query, entityType)); - } else if (queryMethod.isCollectionQuery()) { - if (accessor.getPageable().isUnpaged()) { - int itemCount = (int) elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType()); - query.setPageable(PageRequest.of(0, Math.max(1, itemCount))); - } else { - query.setPageable(accessor.getPageable()); - } - return elasticsearchOperations.queryForList(query, queryMethod.getEntityInformation().getJavaType()); - } else if (tree.isCountProjection()) { - return elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType()); - } + } else if (queryMethod.isCollectionQuery()) { + + if (accessor.getPageable().isUnpaged()) { + + int itemCount = (int) elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType()); + query.setPageable(PageRequest.of(0, Math.max(1, itemCount))); + } else { + + query.setPageable(accessor.getPageable()); + } + + return elasticsearchOperations.queryForList(query, queryMethod.getEntityInformation().getJavaType()); + } else if (tree.isCountProjection()) { + + return elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType()); + } + return elasticsearchOperations.queryForObject(query, queryMethod.getEntityInformation().getJavaType()); } diff --git a/src/test/java/org/springframework/data/elasticsearch/repositories/custommethod/CustomMethodRepositoryBaseTests.java b/src/test/java/org/springframework/data/elasticsearch/repositories/custommethod/CustomMethodRepositoryBaseTests.java index 5f1f686c8..2705908b1 100644 --- a/src/test/java/org/springframework/data/elasticsearch/repositories/custommethod/CustomMethodRepositoryBaseTests.java +++ b/src/test/java/org/springframework/data/elasticsearch/repositories/custommethod/CustomMethodRepositoryBaseTests.java @@ -29,6 +29,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.UUID; import java.util.stream.Stream; import org.junit.Test; @@ -59,6 +60,7 @@ import org.springframework.data.geo.Point; * @author Christoph Strobl * @author Don Wellington * @author Peter-Josef Meisch + * @author Rasmus Faber-Espensen */ public abstract class CustomMethodRepositoryBaseTests { @@ -1239,6 +1241,7 @@ public abstract class CustomMethodRepositoryBaseTests { @Test // DATAES-605 public void streamMethodsShouldWorkWithLargeResultSets() { + // given List entities = createSampleEntities("abc", 10001); repository.saveAll(entities); @@ -1253,6 +1256,7 @@ public abstract class CustomMethodRepositoryBaseTests { @Test // DATAES-605 public void streamMethodsCanHandlePageable() { + // given List entities = createSampleEntities("abc", 10); repository.saveAll(entities); @@ -1271,7 +1275,7 @@ public abstract class CustomMethodRepositoryBaseTests { for (int i = 0; i < numberOfEntities; i++) { SampleEntity entity = new SampleEntity(); - entity.setId(randomNumeric(32)); + entity.setId(UUID.randomUUID().toString()); entity.setAvailable(true); entity.setMessage("Message"); entity.setType(type);