DATAES-631 - Refactored count methods in ElasticsearchTemplate.

Author: sothawo

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

@@ -347,7 +347,9 @@ default List<List<?>> queryForList(List<SearchQuery> queries, List<Class<?>> cla
 	 * @param query
 	 * @return
 	 */
-	<T> long count(CriteriaQuery query);
+	default <T> long count(CriteriaQuery query) {
+		return count(query, null);
+	}
 
 	/**
 	 * return number of elements found by given query
@@ -364,7 +366,9 @@ default List<List<?>> queryForList(List<SearchQuery> queries, List<Class<?>> cla
 	 * @param query
 	 * @return
 	 */
-	<T> long count(SearchQuery query);
+	default <T> long count(SearchQuery query) {
+		return count(query, null);
+	}
 
 	/**
 	 * Execute a multiGet against elasticsearch for the given ids

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

@@ -422,78 +422,50 @@ private <T> CloseableIterator<T> doStream(long scrollTimeInMillis, ScrolledPage<
 
 	@Override
 	public <T> long count(CriteriaQuery criteriaQuery, Class<T> clazz) {
-		QueryBuilder elasticsearchQuery = new CriteriaQueryProcessor().createQueryFromCriteria(criteriaQuery.getCriteria());
-		QueryBuilder elasticsearchFilter = new CriteriaFilterProcessor()
-				.createFilterFromCriteria(criteriaQuery.getCriteria());
-
-		if (elasticsearchFilter == null) {
-			return doCount(prepareCount(criteriaQuery, clazz), elasticsearchQuery);
-		} else {
-			// filter could not be set into CountRequestBuilder, convert request into search request
-			return doCount(prepareSearch(criteriaQuery, clazz), elasticsearchQuery, elasticsearchFilter);
-		}
+		return doCount(criteriaQuery, clazz);
 	}
 
 	@Override
 	public <T> long count(SearchQuery searchQuery, Class<T> clazz) {
-		QueryBuilder elasticsearchQuery = searchQuery.getQuery();
-		QueryBuilder elasticsearchFilter = searchQuery.getFilter();
-
-		if (elasticsearchFilter == null) {
-			return doCount(prepareCount(searchQuery, clazz), elasticsearchQuery);
-		} else {
-			// filter could not be set into CountRequestBuilder, convert request into search request
-			return doCount(prepareSearch(searchQuery, clazz), elasticsearchQuery, elasticsearchFilter);
-		}
+		return doCount(searchQuery, clazz);
 	}
 
-	@Override
-	public <T> long count(CriteriaQuery query) {
-		return count(query, null);
-	}
+	private <T> long doCount(Query query, Class<T> clazz) {
+		SearchRequestBuilder searchRequestBuilder = mapQuery(query, clazz);
+		searchRequestBuilder.setSize(0);
 
-	@Override
-	public <T> long count(SearchQuery query) {
-		return count(query, null);
+		return SearchHitsUtil.getTotalCount(searchRequestBuilder.execute().actionGet().getHits());
 	}
 
-	private long doCount(SearchRequestBuilder countRequestBuilder, QueryBuilder elasticsearchQuery) {
+	private <T> SearchRequestBuilder mapQuery(Query query, Class<T> clazz) {
+		QueryBuilder elasticsearchQuery;
+		QueryBuilder elasticsearchFilter;
 
-		if (elasticsearchQuery != null) {
-			countRequestBuilder.setQuery(elasticsearchQuery);
+		if (query instanceof SearchQuery) {
+			SearchQuery searchQuery = (SearchQuery) query;
+			elasticsearchQuery = searchQuery.getQuery();
+			elasticsearchFilter = searchQuery.getFilter();
+		} else if (query instanceof CriteriaQuery) {
+			CriteriaQuery criteriaQuery = (CriteriaQuery) query;
+			elasticsearchQuery = new CriteriaQueryProcessor().createQueryFromCriteria(criteriaQuery.getCriteria());
+			elasticsearchFilter = new CriteriaFilterProcessor().createFilterFromCriteria(criteriaQuery.getCriteria());
+		} else {
+			throw new IllegalArgumentException("unhandled Query implementation " + query.getClass().getName());
 		}
-		return SearchHitsUtil.getTotalCount(countRequestBuilder.execute().actionGet().getHits());
-	}
 
-	private long doCount(SearchRequestBuilder searchRequestBuilder, QueryBuilder elasticsearchQuery,
-			QueryBuilder elasticsearchFilter) {
+		SearchRequestBuilder searchRequestBuilder = prepareSearch(query, clazz);
+
 		if (elasticsearchQuery != null) {
 			searchRequestBuilder.setQuery(elasticsearchQuery);
 		} else {
 			searchRequestBuilder.setQuery(QueryBuilders.matchAllQuery());
 		}
+
 		if (elasticsearchFilter != null) {
 			searchRequestBuilder.setPostFilter(elasticsearchFilter);
 		}
-		return SearchHitsUtil.getTotalCount(searchRequestBuilder.execute().actionGet().getHits());
-	}
-
-	private <T> SearchRequestBuilder prepareCount(Query query, Class<T> clazz) {
-		String indexName[] = !isEmpty(query.getIndices())
-				? query.getIndices().toArray(new String[query.getIndices().size()])
-				: retrieveIndexNameFromPersistentEntity(clazz);
-		String types[] = !isEmpty(query.getTypes()) ? query.getTypes().toArray(new String[query.getTypes().size()])
-				: retrieveTypeFromPersistentEntity(clazz);
-
-		Assert.notNull(indexName, "No index defined for Query");
-
-		SearchRequestBuilder countRequestBuilder = client.prepareSearch(indexName);
 
-		if (types != null) {
-			countRequestBuilder.setTypes(types);
-		}
-		countRequestBuilder.setSize(0);
-		return countRequestBuilder;
+		return searchRequestBuilder;
 	}
 
 	@Override
@@ -1033,6 +1005,7 @@ private <T> SearchRequestBuilder prepareSearch(Query query, Class<T> clazz) {
 
 	private SearchRequestBuilder prepareSearch(Query query, @Nullable ElasticsearchPersistentEntity<?> entity) {
 		Assert.notNull(query.getIndices(), "No index defined for Query");
+		Assert.notEmpty(query.getIndices(), "No index defined for Query");
 		Assert.notNull(query.getTypes(), "No type defined for Query");
 
 		int startRecord = 0;
@@ -1196,10 +1169,18 @@ private void setPersistentEntityId(Object entity, String id) {
 
 	private void setPersistentEntityIndexAndType(Query query, Class clazz) {
 		if (query.getIndices().isEmpty()) {
-			query.addIndices(retrieveIndexNameFromPersistentEntity(clazz));
+			String[] indices = retrieveIndexNameFromPersistentEntity(clazz);
+
+			if (indices != null) {
+				query.addIndices(indices);
+			}
 		}
 		if (query.getTypes().isEmpty()) {
-			query.addTypes(retrieveTypeFromPersistentEntity(clazz));
+			String[] types = retrieveTypeFromPersistentEntity(clazz);
+
+			if (types != null) {
+				query.addTypes(types);
+			}
 		}
 	}