Skip to content

Commit 1ea9920

Browse files
kevinleturcmohsinh
kevinleturc
authored and
mohsinh
committed
DATAES-106 - Add support for countBy projection
1 parent c866664 commit 1ea9920

File tree

8 files changed

+818
-21
lines changed

8 files changed

+818
-21
lines changed

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
*
3131
* @author Rizwan Idrees
3232
* @author Mohsin Husen
33+
* @author Kevin Leturc
3334
*/
3435
public interface ElasticsearchOperations {
3536

@@ -246,6 +247,23 @@ public interface ElasticsearchOperations {
246247
*/
247248
<T> List<String> queryForIds(SearchQuery query);
248249

250+
/**
251+
* return number of elements found by given query
252+
*
253+
* @param query
254+
* @param clazz
255+
* @return
256+
*/
257+
<T> long count(CriteriaQuery query, Class<T> clazz);
258+
259+
/**
260+
* return number of elements found by given query
261+
*
262+
* @param query
263+
* @return
264+
*/
265+
<T> long count(CriteriaQuery query);
266+
249267
/**
250268
* return number of elements found by given query
251269
*

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

Lines changed: 60 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
import org.elasticsearch.action.mlt.MoreLikeThisRequestBuilder;
5151
import org.elasticsearch.action.search.SearchRequestBuilder;
5252
import org.elasticsearch.action.search.SearchResponse;
53+
import org.elasticsearch.action.search.SearchType;
5354
import org.elasticsearch.action.update.UpdateRequestBuilder;
5455
import org.elasticsearch.action.update.UpdateResponse;
5556
import org.elasticsearch.client.Client;
@@ -96,6 +97,7 @@
9697
* @author Rizwan Idrees
9798
* @author Mohsin Husen
9899
* @author Artur Konczak
100+
* @author Kevin Leturc
99101
*/
100102

101103
public class ElasticsearchTemplate implements ElasticsearchOperations, ApplicationContextAware {
@@ -318,28 +320,75 @@ public <T> FacetedPage<T> queryForPage(StringQuery query, Class<T> clazz, Search
318320
}
319321

320322
@Override
321-
public <T> long count(SearchQuery searchQuery, Class<T> clazz) {
322-
String indexName[] = isNotEmpty(searchQuery.getIndices()) ? searchQuery.getIndices().toArray(new String[searchQuery.getIndices().size()]) : retrieveIndexNameFromPersistentEntity(clazz);
323-
String types[] = isNotEmpty(searchQuery.getTypes()) ? searchQuery.getTypes().toArray(new String[searchQuery.getTypes().size()]) : retrieveTypeFromPersistentEntity(clazz);
323+
public <T> long count(CriteriaQuery criteriaQuery, Class<T> clazz) {
324+
QueryBuilder elasticsearchQuery = new CriteriaQueryProcessor().createQueryFromCriteria(criteriaQuery.getCriteria());
325+
FilterBuilder elasticsearchFilter = new CriteriaFilterProcessor().createFilterFromCriteria(criteriaQuery.getCriteria());
324326

325-
Assert.notNull(indexName, "No index defined for Query");
327+
if (elasticsearchFilter == null) {
328+
return doCount(prepareCount(criteriaQuery, clazz), elasticsearchQuery);
329+
} else {
330+
// filter could not be set into CountRequestBuilder, convert request into search request
331+
return doCount(prepareSearch(criteriaQuery, clazz), elasticsearchQuery, elasticsearchFilter);
332+
}
333+
}
326334

327-
CountRequestBuilder countRequestBuilder = client.prepareCount(indexName);
335+
@Override
336+
public <T> long count(SearchQuery searchQuery, Class<T> clazz) {
337+
QueryBuilder elasticsearchQuery = searchQuery.getQuery();
338+
FilterBuilder elasticsearchFilter = searchQuery.getFilter();
328339

329-
if (types != null) {
330-
countRequestBuilder.setTypes(types);
331-
}
332-
if (searchQuery.getQuery() != null) {
333-
countRequestBuilder.setQuery(searchQuery.getQuery());
340+
if (elasticsearchFilter == null) {
341+
return doCount(prepareCount(searchQuery, clazz), elasticsearchQuery);
342+
} else {
343+
// filter could not be set into CountRequestBuilder, convert request into search request
344+
return doCount(prepareSearch(searchQuery, clazz), elasticsearchQuery, elasticsearchFilter);
334345
}
335-
return countRequestBuilder.execute().actionGet().getCount();
346+
}
347+
348+
@Override
349+
public <T> long count(CriteriaQuery query) {
350+
return count(query, null);
336351
}
337352

338353
@Override
339354
public <T> long count(SearchQuery query) {
340355
return count(query, null);
341356
}
342357

358+
private long doCount(CountRequestBuilder countRequestBuilder, QueryBuilder elasticsearchQuery) {
359+
if (elasticsearchQuery != null) {
360+
countRequestBuilder.setQuery(elasticsearchQuery);
361+
}
362+
return countRequestBuilder.execute().actionGet().getCount();
363+
}
364+
365+
private long doCount(SearchRequestBuilder searchRequestBuilder, QueryBuilder elasticsearchQuery, FilterBuilder elasticsearchFilter) {
366+
if (elasticsearchQuery != null) {
367+
searchRequestBuilder.setQuery(elasticsearchQuery);
368+
} else {
369+
searchRequestBuilder.setQuery(QueryBuilders.matchAllQuery());
370+
}
371+
if (elasticsearchFilter != null) {
372+
searchRequestBuilder.setPostFilter(elasticsearchFilter);
373+
}
374+
searchRequestBuilder.setSearchType(SearchType.COUNT);
375+
return searchRequestBuilder.execute().actionGet().getHits().getTotalHits();
376+
}
377+
378+
private <T> CountRequestBuilder prepareCount(Query query, Class<T> clazz) {
379+
String indexName[] = isNotEmpty(query.getIndices()) ? query.getIndices().toArray(new String[query.getIndices().size()]) : retrieveIndexNameFromPersistentEntity(clazz);
380+
String types[] = isNotEmpty(query.getTypes()) ? query.getTypes().toArray(new String[query.getTypes().size()]) : retrieveTypeFromPersistentEntity(clazz);
381+
382+
Assert.notNull(indexName, "No index defined for Query");
383+
384+
CountRequestBuilder countRequestBuilder = client.prepareCount(indexName);
385+
386+
if (types != null) {
387+
countRequestBuilder.setTypes(types);
388+
}
389+
return countRequestBuilder;
390+
}
391+
343392
@Override
344393
public <T> LinkedList<T> multiGet(SearchQuery searchQuery, Class<T> clazz) {
345394
return resultsMapper.mapResults(getMultiResponse(searchQuery, clazz), clazz);

‎src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchPartQuery.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
*
2929
* @author Rizwan Idrees
3030
* @author Mohsin Husen
31+
* @author Kevin Leturc
3132
*/
3233
public class ElasticsearchPartQuery extends AbstractElasticsearchRepositoryQuery {
3334

@@ -52,6 +53,8 @@ public Object execute(Object[] parameters) {
5253
query.setPageable(accessor.getPageable());
5354
}
5455
return elasticsearchOperations.queryForList(query, queryMethod.getEntityInformation().getJavaType());
56+
} else if (tree.isCountProjection()) {
57+
return elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType());
5558
}
5659
return elasticsearchOperations.queryForObject(query, queryMethod.getEntityInformation().getJavaType());
5760
}

‎src/main/java/org/springframework/data/elasticsearch/repository/query/ElasticsearchQueryMethod.java

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,6 @@
1919

2020
import org.springframework.core.annotation.AnnotationUtils;
2121
import org.springframework.data.elasticsearch.annotations.Query;
22-
import org.springframework.data.elasticsearch.repository.support.ElasticsearchEntityInformation;
23-
import org.springframework.data.elasticsearch.repository.support.ElasticsearchEntityInformationCreator;
2422
import org.springframework.data.repository.core.RepositoryMetadata;
2523
import org.springframework.data.repository.query.QueryMethod;
2624
import org.springframework.util.StringUtils;
@@ -33,14 +31,10 @@
3331
*/
3432
public class ElasticsearchQueryMethod extends QueryMethod {
3533

36-
private final ElasticsearchEntityInformation<?, ?> entityInformation;
3734
private Method method;
3835

39-
public ElasticsearchQueryMethod(Method method, RepositoryMetadata metadata,
40-
ElasticsearchEntityInformationCreator elasticsearchEntityInformationCreator) {
36+
public ElasticsearchQueryMethod(Method method, RepositoryMetadata metadata) {
4137
super(method, metadata);
42-
this.entityInformation = elasticsearchEntityInformationCreator.getEntityInformation(metadata
43-
.getReturnedDomainClass(method));
4438
this.method = method;
4539
}
4640

‎src/main/java/org/springframework/data/elasticsearch/repository/support/ElasticsearchRepositoryFactory.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ private class ElasticsearchQueryLookupStrategy implements QueryLookupStrategy {
105105
@Override
106106
public RepositoryQuery resolveQuery(Method method, RepositoryMetadata metadata, NamedQueries namedQueries) {
107107

108-
ElasticsearchQueryMethod queryMethod = new ElasticsearchQueryMethod(method, metadata, entityInformationCreator);
108+
ElasticsearchQueryMethod queryMethod = new ElasticsearchQueryMethod(method, metadata);
109109
String namedQueryName = queryMethod.getNamedQueryName();
110110

111111
if (namedQueries.hasQuery(namedQueryName)) {

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

Lines changed: 156 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
* @author Mohsin Husen
5959
* @author Franck Marchand
6060
* @author Abdul Mohammed
61+
* @author Kevin Leturc
6162
*/
6263
@RunWith(SpringJUnit4ClassRunner.class)
6364
@ContextConfiguration("classpath:elasticsearch-template-test.xml")
@@ -80,6 +81,26 @@ public void before() {
8081
elasticsearchTemplate.refresh(SampleEntity.class, true);
8182
}
8283

84+
/*
85+
DATAES-106
86+
*/
87+
@Test
88+
public void shouldReturnCountForGivenCriteriaQuery() {
89+
// given
90+
String documentId = randomNumeric(5);
91+
SampleEntity sampleEntity = new SampleEntityBuilder(documentId).message("some message")
92+
.version(System.currentTimeMillis()).build();
93+
94+
IndexQuery indexQuery = getIndexQuery(sampleEntity);
95+
elasticsearchTemplate.index(indexQuery);
96+
elasticsearchTemplate.refresh(SampleEntity.class, true);
97+
CriteriaQuery criteriaQuery = new CriteriaQuery(new Criteria());
98+
// when
99+
long count = elasticsearchTemplate.count(criteriaQuery, SampleEntity.class);
100+
// then
101+
assertThat(count, is(equalTo(1L)));
102+
}
103+
83104
@Test
84105
public void shouldReturnCountForGivenSearchQuery() {
85106
// given
@@ -1187,6 +1208,27 @@ public void shouldIndexSampleEntityWithIndexAndTypeAtRuntime() {
11871208
assertThat(sampleEntities.getTotalElements(), greaterThanOrEqualTo(1L));
11881209
}
11891210

1211+
/*
1212+
DATAES-106
1213+
*/
1214+
@Test
1215+
public void shouldReturnCountForGivenCriteriaQueryWithGivenIndexUsingCriteriaQuery() {
1216+
// given
1217+
String documentId = randomNumeric(5);
1218+
SampleEntity sampleEntity = new SampleEntityBuilder(documentId).message("some message")
1219+
.version(System.currentTimeMillis()).build();
1220+
1221+
IndexQuery indexQuery = getIndexQuery(sampleEntity);
1222+
elasticsearchTemplate.index(indexQuery);
1223+
elasticsearchTemplate.refresh(SampleEntity.class, true);
1224+
CriteriaQuery criteriaQuery = new CriteriaQuery(new Criteria());
1225+
criteriaQuery.addIndices("test-index");
1226+
// when
1227+
long count = elasticsearchTemplate.count(criteriaQuery);
1228+
// then
1229+
assertThat(count, is(equalTo(1L)));
1230+
}
1231+
11901232
/*
11911233
DATAES-67
11921234
*/
@@ -1210,6 +1252,28 @@ public void shouldReturnCountForGivenSearchQueryWithGivenIndexUsingSearchQuery()
12101252
assertThat(count, is(equalTo(1L)));
12111253
}
12121254

1255+
/*
1256+
DATAES-106
1257+
*/
1258+
@Test
1259+
public void shouldReturnCountForGivenCriteriaQueryWithGivenIndexAndTypeUsingCriteriaQuery() {
1260+
// given
1261+
String documentId = randomNumeric(5);
1262+
SampleEntity sampleEntity = new SampleEntityBuilder(documentId).message("some message")
1263+
.version(System.currentTimeMillis()).build();
1264+
1265+
IndexQuery indexQuery = getIndexQuery(sampleEntity);
1266+
elasticsearchTemplate.index(indexQuery);
1267+
elasticsearchTemplate.refresh(SampleEntity.class, true);
1268+
CriteriaQuery criteriaQuery = new CriteriaQuery(new Criteria());
1269+
criteriaQuery.addIndices("test-index");
1270+
criteriaQuery.addTypes("test-type");
1271+
// when
1272+
long count = elasticsearchTemplate.count(criteriaQuery);
1273+
// then
1274+
assertThat(count, is(equalTo(1L)));
1275+
}
1276+
12131277
/*
12141278
DATAES-67
12151279
*/
@@ -1234,6 +1298,43 @@ public void shouldReturnCountForGivenSearchQueryWithGivenIndexAndTypeUsingSearch
12341298
assertThat(count, is(equalTo(1L)));
12351299
}
12361300

1301+
/*
1302+
DATAES-106
1303+
*/
1304+
@Test
1305+
public void shouldReturnCountForGivenCriteriaQueryWithGivenMultiIndices() {
1306+
// given
1307+
cleanUpIndices();
1308+
String documentId1 = randomNumeric(5);
1309+
SampleEntity sampleEntity1 = new SampleEntityBuilder(documentId1).message("some message")
1310+
.version(System.currentTimeMillis()).build();
1311+
1312+
IndexQuery indexQuery1 = new IndexQueryBuilder().withId(sampleEntity1.getId())
1313+
.withIndexName("test-index-1")
1314+
.withObject(sampleEntity1)
1315+
.build();
1316+
1317+
String documentId2 = randomNumeric(5);
1318+
SampleEntity sampleEntity2 = new SampleEntityBuilder(documentId2).message("some test message")
1319+
.version(System.currentTimeMillis()).build();
1320+
1321+
IndexQuery indexQuery2 = new IndexQueryBuilder().withId(sampleEntity2.getId())
1322+
.withIndexName("test-index-2")
1323+
.withObject(sampleEntity2)
1324+
.build();
1325+
1326+
elasticsearchTemplate.bulkIndex(Arrays.asList(indexQuery1, indexQuery2));
1327+
elasticsearchTemplate.refresh("test-index-1", true);
1328+
elasticsearchTemplate.refresh("test-index-2", true);
1329+
1330+
CriteriaQuery criteriaQuery = new CriteriaQuery(new Criteria());
1331+
criteriaQuery.addIndices("test-index-1", "test-index-2");
1332+
// when
1333+
long count = elasticsearchTemplate.count(criteriaQuery);
1334+
// then
1335+
assertThat(count, is(equalTo(2L)));
1336+
}
1337+
12371338
/*
12381339
DATAES-67
12391340
*/
@@ -1310,6 +1411,43 @@ public void shouldDeleteIndexForSpecifiedIndexName() {
13101411
assertThat(elasticsearchTemplate.indexExists("test-index"), is(false));
13111412
}
13121413

1414+
/*
1415+
DATAES-106
1416+
*/
1417+
@Test
1418+
public void shouldReturnCountForGivenCriteriaQueryWithGivenIndexNameForSpecificIndex() {
1419+
// given
1420+
cleanUpIndices();
1421+
String documentId1 = randomNumeric(5);
1422+
SampleEntity sampleEntity1 = new SampleEntityBuilder(documentId1).message("some message")
1423+
.version(System.currentTimeMillis()).build();
1424+
1425+
IndexQuery indexQuery1 = new IndexQueryBuilder().withId(sampleEntity1.getId())
1426+
.withIndexName("test-index-1")
1427+
.withObject(sampleEntity1)
1428+
.build();
1429+
1430+
String documentId2 = randomNumeric(5);
1431+
SampleEntity sampleEntity2 = new SampleEntityBuilder(documentId2).message("some test message")
1432+
.version(System.currentTimeMillis()).build();
1433+
1434+
IndexQuery indexQuery2 = new IndexQueryBuilder().withId(sampleEntity2.getId())
1435+
.withIndexName("test-index-2")
1436+
.withObject(sampleEntity2)
1437+
.build();
1438+
1439+
elasticsearchTemplate.bulkIndex(Arrays.asList(indexQuery1, indexQuery2));
1440+
elasticsearchTemplate.refresh("test-index-1", true);
1441+
elasticsearchTemplate.refresh("test-index-2", true);
1442+
1443+
CriteriaQuery criteriaQuery = new CriteriaQuery(new Criteria());
1444+
criteriaQuery.addIndices("test-index-1");
1445+
// when
1446+
long count = elasticsearchTemplate.count(criteriaQuery);
1447+
// then
1448+
assertThat(count, is(equalTo(1L)));
1449+
}
1450+
13131451
/*
13141452
DATAES-67
13151453
*/
@@ -1349,11 +1487,28 @@ public void shouldReturnCountForGivenSearchQueryWithGivenIndexNameForSpecificInd
13491487
assertThat(count, is(equalTo(1L)));
13501488
}
13511489

1490+
@Test(expected = IllegalArgumentException.class)
1491+
public void shouldThrowAnExceptionForGivenCriteriaQueryWhenNoIndexSpecifiedForCountQuery() {
1492+
// given
1493+
String documentId = randomNumeric(5);
1494+
SampleEntity sampleEntity = new SampleEntityBuilder(documentId).message("some message")
1495+
.version(System.currentTimeMillis()).build();
1496+
1497+
IndexQuery indexQuery = getIndexQuery(sampleEntity);
1498+
elasticsearchTemplate.index(indexQuery);
1499+
elasticsearchTemplate.refresh(SampleEntity.class, true);
1500+
CriteriaQuery criteriaQuery = new CriteriaQuery(new Criteria());
1501+
// when
1502+
long count = elasticsearchTemplate.count(criteriaQuery);
1503+
// then
1504+
assertThat(count, is(equalTo(1L)));
1505+
}
1506+
13521507
/*
13531508
DATAES-67
13541509
*/
13551510
@Test(expected = IllegalArgumentException.class)
1356-
public void shouldThrowAnExceptionWhenNoIndexSpecifiedForCountQuery() {
1511+
public void shouldThrowAnExceptionForGivenSearchQueryWhenNoIndexSpecifiedForCountQuery() {
13571512
// given
13581513
String documentId = randomNumeric(5);
13591514
SampleEntity sampleEntity = new SampleEntityBuilder(documentId).message("some message")

0 commit comments

Comments
 (0)