From 47d3277bd5abe94e813b0aeb592daba7df2b4527 Mon Sep 17 00:00:00 2001 From: Mohsin Husen Date: Tue, 18 Mar 2014 19:46:56 +0000 Subject: [PATCH] DATAES-67 - count request in ElasticsearchTemplate does not respect specified index --- .../core/ElasticsearchOperations.java | 12 +++- .../core/ElasticsearchTemplate.java | 36 +++++++--- .../core/ElasticsearchTemplateTests.java | 69 +++++++++++++++++++ 3 files changed, 106 insertions(+), 11 deletions(-) diff --git a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java index ac581b959..51bcbe77e 100755 --- a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.java +++ b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchOperations.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. @@ -174,7 +174,7 @@ public interface ElasticsearchOperations { List queryForIds(SearchQuery query); /** - * return number of elements found by for given query + * return number of elements found by given query * * @param query * @param clazz @@ -182,6 +182,14 @@ public interface ElasticsearchOperations { */ long count(SearchQuery query, Class clazz); + /** + * return number of elements found by given query + * + * @param query + * @return + */ + long count(SearchQuery query); + /** * Execute a multiGet against elasticsearch for the given ids * diff --git a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java index e41b906d9..816327e63 100755 --- a/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.java +++ b/src/main/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplate.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. @@ -239,16 +239,28 @@ public class ElasticsearchTemplate implements ElasticsearchOperations { } @Override - public long count(SearchQuery query, Class clazz) { - ElasticsearchPersistentEntity persistentEntity = getPersistentEntityFor(clazz); - CountRequestBuilder countRequestBuilder = client.prepareCount(persistentEntity.getIndexName()).setTypes( - persistentEntity.getIndexType()); - if (query.getQuery() != null) { - countRequestBuilder.setQuery(query.getQuery()); + public long count(SearchQuery searchQuery, Class clazz) { + String indexName[] = isNotEmpty(searchQuery.getIndices()) ? searchQuery.getIndices().toArray(new String[searchQuery.getIndices().size()]) : retrieveIndexNameFromPersistentEntity(clazz); + String types[] = isNotEmpty(searchQuery.getTypes()) ? searchQuery.getTypes().toArray(new String[searchQuery.getTypes().size()]) : retrieveTypeFromPersistentEntity(clazz); + + Assert.notNull(indexName, "No index defined for Query"); + + CountRequestBuilder countRequestBuilder = client.prepareCount(indexName); + + if (types != null) { + countRequestBuilder.setTypes(types); + } + if (searchQuery.getQuery() != null) { + countRequestBuilder.setQuery(searchQuery.getQuery()); } return countRequestBuilder.execute().actionGet().getCount(); } + @Override + public long count(SearchQuery query) { + return count(query, null); + } + @Override public LinkedList multiGet(SearchQuery searchQuery, Class clazz) { return resultsMapper.mapResults(getMultiResponse(searchQuery, clazz), clazz); @@ -704,11 +716,17 @@ public class ElasticsearchTemplate implements ElasticsearchOperations { } private String[] retrieveIndexNameFromPersistentEntity(Class clazz) { - return new String[]{getPersistentEntityFor(clazz).getIndexName()}; + if (clazz != null) { + return new String[]{getPersistentEntityFor(clazz).getIndexName()}; + } + return null; } private String[] retrieveTypeFromPersistentEntity(Class clazz) { - return new String[]{getPersistentEntityFor(clazz).getIndexType()}; + if (clazz != null) { + return new String[]{getPersistentEntityFor(clazz).getIndexType()}; + } + return null; } private List extractIds(SearchResponse response) { diff --git a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java b/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java index e8283c4ea..5feeafc3e 100755 --- a/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java +++ b/src/test/java/org/springframework/data/elasticsearch/core/ElasticsearchTemplateTests.java @@ -1167,6 +1167,75 @@ public class ElasticsearchTemplateTests { assertThat(sampleEntities.getTotalElements(), greaterThanOrEqualTo(1L)); } + /* + DATAES-67 + */ + @Test + public void shouldReturnCountForGivenSearchQueryWithGivenIndexUsingSearchQuery() { + // given + String documentId = randomNumeric(5); + SampleEntity sampleEntity = new SampleEntityBuilder(documentId).message("some message") + .version(System.currentTimeMillis()).build(); + + IndexQuery indexQuery = getIndexQuery(sampleEntity); + elasticsearchTemplate.index(indexQuery); + elasticsearchTemplate.refresh(SampleEntity.class, true); + SearchQuery searchQuery = new NativeSearchQueryBuilder() + .withQuery(matchAllQuery()) + .withIndices("test-index") + .build(); + // when + long count = elasticsearchTemplate.count(searchQuery); + // then + assertThat(count, is(equalTo(1L))); + } + + /* + DATAES-67 + */ + @Test + public void shouldReturnCountForGivenSearchQueryWithGivenIndexAndTypeUsingSearchQuery() { + // given + String documentId = randomNumeric(5); + SampleEntity sampleEntity = new SampleEntityBuilder(documentId).message("some message") + .version(System.currentTimeMillis()).build(); + + IndexQuery indexQuery = getIndexQuery(sampleEntity); + elasticsearchTemplate.index(indexQuery); + elasticsearchTemplate.refresh(SampleEntity.class, true); + SearchQuery searchQuery = new NativeSearchQueryBuilder() + .withQuery(matchAllQuery()) + .withIndices("test-index") + .withTypes("test-type") + .build(); + // when + long count = elasticsearchTemplate.count(searchQuery); + // then + assertThat(count, is(equalTo(1L))); + } + + /* + DATAES-67 + */ + @Test(expected = IllegalArgumentException.class) + public void shouldThrowAnExceptionWhenNoIndexSpecifiedForCountQuery() { + // given + String documentId = randomNumeric(5); + SampleEntity sampleEntity = new SampleEntityBuilder(documentId).message("some message") + .version(System.currentTimeMillis()).build(); + + IndexQuery indexQuery = getIndexQuery(sampleEntity); + elasticsearchTemplate.index(indexQuery); + elasticsearchTemplate.refresh(SampleEntity.class, true); + SearchQuery searchQuery = new NativeSearchQueryBuilder() + .withQuery(matchAllQuery()) + .build(); + // when + long count = elasticsearchTemplate.count(searchQuery); + // then + assertThat(count, is(equalTo(1L))); + } + private IndexQuery getIndexQuery(SampleEntity sampleEntity) { return new IndexQueryBuilder().withId(sampleEntity.getId()).withObject(sampleEntity).build(); }