HSEARCH-3351 Define a proper ES work for count

hibernate · Dec 12, 2018 · f152378 · f152378
1 parent 6f37f85
commit f152378
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 7 deletions.
diff --git a/...rg/hibernate/search/backend/elasticsearch/search/query/impl/ElasticsearchSearchQuery.java b/...rg/hibernate/search/backend/elasticsearch/search/query/impl/ElasticsearchSearchQuery.java
@@ -6,8 +6,10 @@
  */
 package org.hibernate.search.backend.elasticsearch.search.query.impl;
 
+import java.util.Optional;
 import java.util.Set;
 
+import org.hibernate.search.backend.elasticsearch.gson.impl.JsonAccessor;
 import org.hibernate.search.backend.elasticsearch.util.impl.URLEncodedString;
 import org.hibernate.search.backend.elasticsearch.orchestration.impl.ElasticsearchWorkOrchestrator;
 import org.hibernate.search.backend.elasticsearch.work.impl.ElasticsearchWork;
@@ -77,11 +79,13 @@ public SearchResult<T> execute() {
 
 	@Override
 	public long executeCount() {
-		ElasticsearchWork<SearchResult<T>> work = workFactory.search(
-				indexNames, routingKeys,
-				payload, searchResultExtractor,
-				firstResultIndex, 0L );
-		SearchResult<T> executeNoHits = queryOrchestrator.submit( work ).join();
-		return executeNoHits.getHitCount();
+		JsonObject filteredPayload = new JsonObject();
+		Optional<JsonObject> querySubTree = JsonAccessor.root().property( "query" ).asObject().get( payload );
+		if ( querySubTree.isPresent() ) {
+			filteredPayload.add( "query", querySubTree.get() );
+		}
+
+		ElasticsearchWork<Long> work = workFactory.count( indexNames, routingKeys, filteredPayload );
+		return queryOrchestrator.submit( work ).join();
 	}
 }
diff --git a/...main/java/org/hibernate/search/backend/elasticsearch/work/impl/ElasticsearchStubWork.java b/...main/java/org/hibernate/search/backend/elasticsearch/work/impl/ElasticsearchStubWork.java
@@ -11,6 +11,7 @@
 
 import org.hibernate.search.backend.elasticsearch.client.impl.ElasticsearchRequest;
 import org.hibernate.search.backend.elasticsearch.client.impl.ElasticsearchResponse;
+import org.hibernate.search.backend.elasticsearch.gson.impl.JsonAccessor;
 
 import com.google.gson.JsonObject;
 
@@ -32,6 +33,11 @@ public ElasticsearchStubWork(ElasticsearchRequest request, Function<JsonObject,
 		this.resultFunction = resultFunction;
 	}
 
+	public ElasticsearchStubWork(ElasticsearchRequest request, JsonAccessor<T> accessor) {
+		this.request = request;
+		this.resultFunction = response -> accessor.get( response ).get();
+	}
+
 	@Override
 	public CompletableFuture<T> execute(ElasticsearchWorkExecutionContext context) {
 		CompletableFuture<ElasticsearchResponse> response = context.getClient().submit( request );

diff --git a/...va/org/hibernate/search/backend/elasticsearch/work/impl/ElasticsearchStubWorkFactory.java b/...va/org/hibernate/search/backend/elasticsearch/work/impl/ElasticsearchStubWorkFactory.java
@@ -11,6 +11,7 @@
 
 import org.hibernate.search.backend.elasticsearch.client.impl.ElasticsearchRequest;
 import org.hibernate.search.backend.elasticsearch.client.impl.Paths;
+import org.hibernate.search.backend.elasticsearch.gson.impl.JsonAccessor;
 import org.hibernate.search.backend.elasticsearch.index.settings.impl.esnative.IndexSettings;
 import org.hibernate.search.backend.elasticsearch.multitenancy.impl.MultiTenancyStrategy;
 import org.hibernate.search.backend.elasticsearch.util.impl.URLEncodedString;
@@ -172,4 +173,18 @@ public <T> ElasticsearchWork<SearchResult<T>> search(Set<URLEncodedString> index
 
 		return new ElasticsearchStubWork<>( builder.build(), searchResultExtractor::extract );
 	}
+
+	@Override
+	public ElasticsearchWork<Long> count(Set<URLEncodedString> indexNames, Set<String> routingKeys, JsonObject payload) {
+		ElasticsearchRequest.Builder builder = ElasticsearchRequest.post()
+				.multiValuedPathComponent( indexNames )
+				.pathComponent( Paths._COUNT )
+				.body( payload );
+
+		if ( !routingKeys.isEmpty() ) {
+			builder.param( "_routing", routingKeys.stream().collect( Collectors.joining( "," ) ) );
+		}
+
+		return new ElasticsearchStubWork<>( builder.build(), JsonAccessor.root().property( "count" ).asLong() );
+	}
 }
diff --git a/...n/java/org/hibernate/search/backend/elasticsearch/work/impl/ElasticsearchWorkFactory.java b/...n/java/org/hibernate/search/backend/elasticsearch/work/impl/ElasticsearchWorkFactory.java
@@ -42,4 +42,6 @@ <T> ElasticsearchWork<SearchResult<T>> search(Set<URLEncodedString> indexNames,
 			JsonObject payload, ElasticsearchSearchResultExtractor<T> searchResultExtractor,
 			Long offset, Long limit);
 
+	ElasticsearchWork<Long> count(Set<URLEncodedString> indexNames, Set<String> routingKeys, JsonObject payload);
+
 }
diff --git a/...ernate/search/integrationtest/backend/tck/search/SearchResultLoadingOrTransformingIT.java b/...ernate/search/integrationtest/backend/tck/search/SearchResultLoadingOrTransformingIT.java
@@ -327,7 +327,7 @@ public void projections_hitTransformer_referencesTransformer_objectLoading() {
 	}
 
 	@Test
-	public void count() {
+	public void countQuery() {
 		StubMappingSearchTarget searchTarget = indexManager.createSearchTarget();
 
 		SearchQuery<DocumentReference> query = searchTarget.query()
@@ -345,6 +345,25 @@ public void count() {
 		assertEquals( 1L, query.executeCount() );
 	}
 
+	@Test
+	public void countQueryWithProjection() {
+		StubMappingSearchTarget searchTarget = indexManager.createSearchTarget();
+
+		SearchQuery<String> query = searchTarget.query()
+				.asProjection( f -> f.field( "string", String.class ).toProjection() )
+				.predicate( f -> f.matchAll().toPredicate() )
+				.build();
+
+		assertEquals( 2L, query.executeCount() );
+
+		query = searchTarget.query()
+				.asProjection( f -> f.field( "string", String.class ).toProjection() )
+				.predicate( f -> f.match().onField( "string" ).matching( STRING_VALUE ).toPredicate() )
+				.build();
+
+		assertEquals( 1L, query.executeCount() );
+	}
+
 	private void initData() {
 		IndexWorkPlan<? extends DocumentElement> workPlan = indexManager.createWorkPlan();
 		workPlan.add( referenceProvider( MAIN_ID ), document -> {