DATAES-631 - WIP.

Author: sothawo

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

@@ -532,21 +532,34 @@ public <T> long count(SearchQuery searchQuery, Class<T> clazz) {
 		QueryBuilder elasticsearchFilter = searchQuery.getFilter();
 
 		if (elasticsearchFilter == null) {
-			return doCount(prepareCount(searchQuery, clazz), elasticsearchQuery);
+//			return doCount(prepareCount(searchQuery, clazz),
+//			elasticsearchQuery);
+			return doCount(prepareSearch(searchQuery, clazz), elasticsearchQuery, elasticsearchFilter);
 		} else {
 			// filter could not be set into CountRequestBuilder, convert request into search request
 			return doCount(prepareSearch(searchQuery, clazz), elasticsearchQuery, elasticsearchFilter);
 		}
 	}
 
-	@Override
-	public <T> long count(CriteriaQuery query) {
-		return count(query, null);
-	}
+	private <T> SearchRequest mapQuery(Query query, Class<T> clazz) {
+		QueryBuilder elasticsearchQuery;
+		QueryBuilder elasticsearchFilter;
 
-	@Override
-	public <T> long count(SearchQuery query) {
-		return count(query, null);
+		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());
+		}
+
+		SearchRequest searchRequest = prepareSearch(query, clazz);
+
+		return searchRequest;
 	}
 
 	private long doCount(SearchRequest countRequest, QueryBuilder elasticsearchQuery) {
@@ -572,13 +585,12 @@ private long doCount(SearchRequest searchRequest, QueryBuilder elasticsearchQuer
 		if (elasticsearchFilter != null) {
 			searchRequest.source().postFilter(elasticsearchFilter);
 		}
-		SearchResponse response;
+
 		try {
-			response = client.search(searchRequest, RequestOptions.DEFAULT);
+			return SearchHitsUtil.getTotalCount(client.search(searchRequest, RequestOptions.DEFAULT).getHits());
 		} catch (IOException e) {
 			throw new ElasticsearchException("Error for search request: " + searchRequest.toString(), e);
 		}
-		return SearchHitsUtil.getTotalCount(response.getHits());
 	}
 
 	private <T> SearchRequest prepareCount(Query query, Class<T> clazz) {
@@ -1473,10 +1485,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);
+			}
 		}
 	}