DATAES-907 - Track Total Hits not working when set to false.

Original PR: #515

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

@@ -238,7 +238,7 @@ public abstract class AbstractElasticsearchTemplate implements ElasticsearchOper
 
 	@Override
 	public void delete(Query query, Class<?> clazz) {
-		delete(query, getIndexCoordinatesFor(clazz));
+		delete(query, clazz, getIndexCoordinatesFor(clazz));
 	}
 
 	@Override

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

@@ -256,7 +256,7 @@ public class ElasticsearchRestTemplate extends AbstractElasticsearchTemplate {
 		Assert.notNull(query, "query must not be null");
 		Assert.notNull(index, "index must not be null");
 
-		final boolean trackTotalHits = query.getTrackTotalHits();
+		final Boolean trackTotalHits = query.getTrackTotalHits();
 		query.setTrackTotalHits(true);
 		SearchRequest searchRequest = requestFactory.searchRequest(query, clazz, index);
 		query.setTrackTotalHits(trackTotalHits);

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

@@ -276,7 +276,7 @@ public class ElasticsearchTemplate extends AbstractElasticsearchTemplate {
 		Assert.notNull(query, "query must not be null");
 		Assert.notNull(index, "index must not be null");
 
-		final boolean trackTotalHits = query.getTrackTotalHits();
+		final Boolean trackTotalHits = query.getTrackTotalHits();
 		query.setTrackTotalHits(true);
 		SearchRequestBuilder searchRequestBuilder = requestFactory.searchRequestBuilder(client, query, clazz, index);
 		query.setTrackTotalHits(trackTotalHits);

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

@@ -1152,8 +1152,10 @@ class RequestFactory {
 
 		}
 
-		if (query.getTrackTotalHits()) {
+		if (query.getTrackTotalHits() != null) {
 			sourceBuilder.trackTotalHits(query.getTrackTotalHits());
+		} else if (query.getTrackTotalHitsUpTo() != null) {
+			sourceBuilder.trackTotalHitsUpTo(query.getTrackTotalHitsUpTo());
 		}
 
 		if (StringUtils.hasLength(query.getRoute())) {
@@ -1225,8 +1227,10 @@ class RequestFactory {
 			prepareNativeSearch(searchRequestBuilder, (NativeSearchQuery) query);
 		}
 
-		if (query.getTrackTotalHits()) {
+		if (query.getTrackTotalHits() != null) {
 			searchRequestBuilder.setTrackTotalHits(query.getTrackTotalHits());
+		} else if (query.getTrackTotalHitsUpTo() != null) {
+			searchRequestBuilder.setTrackTotalHitsUpTo(query.getTrackTotalHitsUpTo());
 		}
 
 		if (StringUtils.hasLength(query.getRoute())) {

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

@@ -26,5 +26,9 @@ package org.springframework.data.elasticsearch.core;
  */
 public enum TotalHitsRelation {
 	EQUAL_TO, //
-	GREATER_THAN_OR_EQUAL_TO
+	GREATER_THAN_OR_EQUAL_TO, //
+	/**
+	 * @since 4.1
+	 */
+	OFF
 }

src/main/java/org/springframework/data/elasticsearch/core/document/SearchDocumentResponse.java

@@ -36,9 +36,9 @@ import org.springframework.util.Assert;
  */
 public class SearchDocumentResponse {
 
-	private long totalHits;
+	private final long totalHits;
-	private String totalHitsRelation;
+	private final String totalHitsRelation;
-	private float maxScore;
+	private final float maxScore;
 	private final String scrollId;
 	private final List<SearchDocument> searchDocuments;
 	private final Aggregations aggregations;
@@ -108,8 +108,17 @@ public class SearchDocumentResponse {
 	public static SearchDocumentResponse from(SearchHits searchHits, @Nullable String scrollId,
 			@Nullable Aggregations aggregations) {
 		TotalHits responseTotalHits = searchHits.getTotalHits();
-		long totalHits = responseTotalHits.value;
+
-		String totalHitsRelation = responseTotalHits.relation.name();
+		long totalHits;
+		String totalHitsRelation;
+
+		if (responseTotalHits != null) {
+			totalHits = responseTotalHits.value;
+			totalHitsRelation = responseTotalHits.relation.name();
+		} else {
+			totalHits = searchHits.getHits().length;
+			totalHitsRelation = "OFF";
+		}
 
 		float maxScore = searchHits.getMaxScore();
 

src/main/java/org/springframework/data/elasticsearch/core/query/AbstractQuery.java

@@ -56,7 +56,8 @@ abstract class AbstractQuery implements Query {
 	@Nullable protected String preference;
 	@Nullable protected Integer maxResults;
 	@Nullable protected HighlightQuery highlightQuery;
-	private boolean trackTotalHits = false;
+	@Nullable private Boolean trackTotalHits;
+	@Nullable private Integer trackTotalHitsUpTo;
 	@Nullable private Duration scrollTime;
 
 	@Override
@@ -220,15 +221,27 @@ abstract class AbstractQuery implements Query {
 	}
 
 	@Override
-	public void setTrackTotalHits(boolean trackTotalHits) {
+	public void setTrackTotalHits(@Nullable Boolean trackTotalHits) {
 		this.trackTotalHits = trackTotalHits;
 	}
 
 	@Override
-	public boolean getTrackTotalHits() {
+	@Nullable
+	public Boolean getTrackTotalHits() {
 		return trackTotalHits;
 	}
 
+	@Override
+	public void setTrackTotalHitsUpTo(@Nullable Integer trackTotalHitsUpTo) {
+		this.trackTotalHitsUpTo = trackTotalHitsUpTo;
+	}
+
+	@Override
+	@Nullable
+	public Integer getTrackTotalHitsUpTo() {
+		return trackTotalHitsUpTo;
+	}
+
 	@Nullable
 	@Override
 	public Duration getScrollTime() {

src/main/java/org/springframework/data/elasticsearch/core/query/Query.java

@@ -219,7 +219,7 @@ public interface Query {
 	 * @param trackTotalHits the value to set.
 	 * @since 4.0
 	 */
-	void setTrackTotalHits(boolean trackTotalHits);
+	void setTrackTotalHits(@Nullable Boolean trackTotalHits);
 
 	/**
 	 * Sets the flag whether to set the Track_total_hits parameter on queries {@see <a href=
@@ -229,7 +229,25 @@ public interface Query {
 	 * @return the set value.
 	 * @since 4.0
 	 */
-	boolean getTrackTotalHits();
+	@Nullable
+	Boolean getTrackTotalHits();
+
+	/**
+	 * Sets the maximum value up to which total hits are tracked. Only relevant if #getTrackTotalHits is {@literal null}
+	 * 
+	 * @param trackTotalHitsUpTo max limit for trackTotalHits
+	 * @since 4.1
+	 */
+	void setTrackTotalHitsUpTo(@Nullable Integer trackTotalHitsUpTo);
+
+	/**
+	 * Gets the maximum value up to which total hits are tracked. Only relevant if #getTrackTotalHits is {@literal null}
+	 *
+	 * @return max limit for trackTotalHits
+	 * @since 4.1
+	 */
+	@Nullable
+	Integer getTrackTotalHitsUpTo();
 
 	/**
 	 * For queries that are used in delete request, these are internally handled by Elasticsearch as scroll/bulk delete

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

@@ -32,6 +32,7 @@ import org.elasticsearch.action.support.ActiveShardCount;
 import org.elasticsearch.action.support.WriteRequest;
 import org.elasticsearch.action.update.UpdateRequest;
 import org.elasticsearch.common.unit.TimeValue;
+import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 import org.springframework.data.annotation.Id;
 import org.springframework.data.elasticsearch.UncategorizedElasticsearchException;
@@ -59,6 +60,7 @@ import org.springframework.test.context.ContextConfiguration;
  */
 @SpringIntegrationTest
 @ContextConfiguration(classes = { ElasticsearchRestTemplateConfiguration.class })
+@DisplayName("ElasticsearchRestTemplate")
 public class ElasticsearchRestTemplateTests extends ElasticsearchTemplateTests {
 
 	@Test

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

@@ -44,7 +44,9 @@ import java.util.Set;
 import java.util.UUID;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
+import org.assertj.core.api.SoftAssertions;
 import org.assertj.core.util.Lists;
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.support.IndicesOptions;
@@ -61,6 +63,7 @@ import org.elasticsearch.search.sort.SortBuilders;
 import org.elasticsearch.search.sort.SortOrder;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.dao.OptimisticLockingFailureException;
@@ -3511,6 +3514,97 @@ public abstract class ElasticsearchTemplateTests {
 		assertThatSeqNoPrimaryTermIsFilled(entity2);
 	}
 
+	@Test // DATAES-907
+	@DisplayName("should track_total_hits with default value")
+	void shouldTrackTotalHitsWithDefaultValue() {
+
+		NativeSearchQuery queryAll = new NativeSearchQueryBuilder().withQuery(matchAllQuery()).build();
+		operations.delete(queryAll, SampleEntity.class);
+
+		List<SampleEntity> entities = IntStream.rangeClosed(1, 15_000)
+				.mapToObj(i -> SampleEntity.builder().id("" + i).build()).collect(Collectors.toList());
+
+		operations.save(entities);
+		indexOperations.refresh();
+
+		queryAll.setTrackTotalHits(null);
+		SearchHits<SampleEntity> searchHits = operations.search(queryAll, SampleEntity.class);
+
+		SoftAssertions softly = new SoftAssertions();
+		softly.assertThat(searchHits.getTotalHits()).isEqualTo((long) RequestFactory.INDEX_MAX_RESULT_WINDOW);
+		softly.assertThat(searchHits.getTotalHitsRelation()).isEqualTo(TotalHitsRelation.GREATER_THAN_OR_EQUAL_TO);
+		softly.assertAll();
+	}
+
+	@Test // DATAES-907
+	@DisplayName("should track total hits")
+	void shouldTrackTotalHits() {
+
+		NativeSearchQuery queryAll = new NativeSearchQueryBuilder().withQuery(matchAllQuery()).build();
+		operations.delete(queryAll, SampleEntity.class);
+
+		List<SampleEntity> entities = IntStream.rangeClosed(1, 15_000)
+				.mapToObj(i -> SampleEntity.builder().id("" + i).build()).collect(Collectors.toList());
+
+		operations.save(entities);
+		indexOperations.refresh();
+
+		queryAll.setTrackTotalHits(true);
+		queryAll.setTrackTotalHitsUpTo(12_345);
+		SearchHits<SampleEntity> searchHits = operations.search(queryAll, SampleEntity.class);
+
+		SoftAssertions softly = new SoftAssertions();
+		softly.assertThat(searchHits.getTotalHits()).isEqualTo(15_000L);
+		softly.assertThat(searchHits.getTotalHitsRelation()).isEqualTo(TotalHitsRelation.EQUAL_TO);
+		softly.assertAll();
+	}
+
+	@Test // DATAES-907
+	@DisplayName("should track total hits to specific value")
+	void shouldTrackTotalHitsToSpecificValue() {
+
+		NativeSearchQuery queryAll = new NativeSearchQueryBuilder().withQuery(matchAllQuery()).build();
+		operations.delete(queryAll, SampleEntity.class);
+
+		List<SampleEntity> entities = IntStream.rangeClosed(1, 15_000)
+				.mapToObj(i -> SampleEntity.builder().id("" + i).build()).collect(Collectors.toList());
+
+		operations.save(entities);
+		indexOperations.refresh();
+
+		queryAll.setTrackTotalHits(null);
+		queryAll.setTrackTotalHitsUpTo(12_345);
+		SearchHits<SampleEntity> searchHits = operations.search(queryAll, SampleEntity.class);
+
+		SoftAssertions softly = new SoftAssertions();
+		softly.assertThat(searchHits.getTotalHits()).isEqualTo(12_345L);
+		softly.assertThat(searchHits.getTotalHitsRelation()).isEqualTo(TotalHitsRelation.GREATER_THAN_OR_EQUAL_TO);
+		softly.assertAll();
+	}
+
+	@Test
+	@DisplayName("should track total hits is off")
+	void shouldTrackTotalHitsIsOff() {
+
+		NativeSearchQuery queryAll = new NativeSearchQueryBuilder().withQuery(matchAllQuery()).build();
+		operations.delete(queryAll, SampleEntity.class);
+
+		List<SampleEntity> entities = IntStream.rangeClosed(1, 15_000)
+				.mapToObj(i -> SampleEntity.builder().id("" + i).build()).collect(Collectors.toList());
+
+		operations.save(entities);
+		indexOperations.refresh();
+
+		queryAll.setTrackTotalHits(false);
+		queryAll.setTrackTotalHitsUpTo(12_345);
+		SearchHits<SampleEntity> searchHits = operations.search(queryAll, SampleEntity.class);
+
+		SoftAssertions softly = new SoftAssertions();
+		softly.assertThat(searchHits.getTotalHits()).isEqualTo(10_000L);
+		softly.assertThat(searchHits.getTotalHitsRelation()).isEqualTo(TotalHitsRelation.OFF);
+		softly.assertAll();
+	}
+
 	@Data
 	@NoArgsConstructor
 	@AllArgsConstructor

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

@@ -34,6 +34,7 @@ import org.elasticsearch.action.update.UpdateRequestBuilder;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.index.engine.DocumentMissingException;
+import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.data.annotation.Id;
@@ -56,6 +57,7 @@ import org.springframework.test.context.ContextConfiguration;
  */
 @SpringIntegrationTest
 @ContextConfiguration(classes = { ElasticsearchTemplateConfiguration.class })
+@DisplayName("ElasticsearchTransportTemplate")
 public class ElasticsearchTransportTemplateTests extends ElasticsearchTemplateTests {
 
 	@Autowired private Client client;