DATAES-67 - count request in ElasticsearchTemplate does not respect specified index

Author: mohsinh

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,14 +174,22 @@ public interface ElasticsearchOperations {
 	<T> List<String> queryForIds(SearchQuery query);
 
 	/**
-	 * return number of elements found by for given query
+	 * return number of elements found by given query
 	 *
 	 * @param query
 	 * @param clazz
 	 * @return
 	 */
 	<T> long count(SearchQuery query, Class<T> clazz);
 
+	/**
+	 * return number of elements found by given query
+	 *
+	 * @param query
+	 * @return
+	 */
+	<T> long count(SearchQuery query);
+
 	/**
 	 * Execute a multiGet against elasticsearch for the given ids
 	 *

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 <T> FacetedPage<T> queryForPage(StringQuery query, Class<T> clazz, Search
 	}
 
 	@Override
-	public <T> long count(SearchQuery query, Class<T> clazz) {
-		ElasticsearchPersistentEntity<T> persistentEntity = getPersistentEntityFor(clazz);
-		CountRequestBuilder countRequestBuilder = client.prepareCount(persistentEntity.getIndexName()).setTypes(
-				persistentEntity.getIndexType());
-		if (query.getQuery() != null) {
-			countRequestBuilder.setQuery(query.getQuery());
+	public <T> long count(SearchQuery searchQuery, Class<T> 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 <T> long count(SearchQuery query) {
+		return count(query, null);
+	}
+
 	@Override
 	public <T> LinkedList<T> multiGet(SearchQuery searchQuery, Class<T> clazz) {
 		return resultsMapper.mapResults(getMultiResponse(searchQuery, clazz), clazz);
@@ -704,11 +716,17 @@ private void setPersistentEntityId(Object entity, String id) {
 	}
 
 	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<String> extractIds(SearchResponse response) {

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

@@ -1167,6 +1167,75 @@ public void shouldIndexSampleEntityWithIndexAndTypeAtRuntime() {
 		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();
 	}