Support _total and _count works for hibernate search (#3567)

Add support for _total, _count, and _offset to Lucene backend.

Co-authored-by: Michael Buckley <michael.buckley@smilecdr.com>
Co-authored-by: juan.marchionatto <juan.marchionatto@smilecdr.com>

hapi-fhir-base/src/main/java/ca/uhn/fhir/i18n/Msg.java

@@ -25,7 +25,7 @@ public final class Msg {
 
 	/**
 	 * IMPORTANT: Please update the following comment after you add a new code
-	 * Last code value: 2076
+	 * Last code value: 2078
 	 */
 
 	private Msg() {}

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/dao/FulltextSearchSvcImpl.java

@@ -36,6 +36,8 @@ import ca.uhn.fhir.jpa.model.search.ExtendedLuceneIndexData;
 import ca.uhn.fhir.jpa.search.autocomplete.ValueSetAutocompleteOptions;
 import ca.uhn.fhir.jpa.search.autocomplete.ValueSetAutocompleteSearch;
 import ca.uhn.fhir.jpa.search.builder.ISearchQueryExecutor;
+import ca.uhn.fhir.jpa.search.builder.SearchQueryExecutors;
+import ca.uhn.fhir.jpa.search.builder.sql.SearchQueryExecutor;
 import ca.uhn.fhir.jpa.searchparam.SearchParameterMap;
 import ca.uhn.fhir.jpa.searchparam.extractor.ISearchParamExtractor;
 import ca.uhn.fhir.jpa.searchparam.extractor.ResourceIndexedSearchParams;
@@ -47,7 +49,9 @@ import ca.uhn.fhir.rest.server.util.ISearchParamRegistry;
 import ca.uhn.fhir.rest.server.util.ResourceSearchParams;
 import org.hibernate.search.backend.elasticsearch.ElasticsearchExtension;
 import org.hibernate.search.engine.search.query.SearchScroll;
+import org.hibernate.search.engine.search.query.dsl.SearchQueryOptionsStep;
 import org.hibernate.search.mapper.orm.Search;
+import org.hibernate.search.mapper.orm.search.loading.dsl.SearchLoadingOptionsStep;
 import org.hibernate.search.mapper.orm.session.SearchSession;
 import org.hibernate.search.mapper.orm.work.SearchIndexingPlan;
 import org.hibernate.search.util.common.SearchException;
@@ -132,14 +136,34 @@ public class FulltextSearchSvcImpl implements IFulltextSearchSvc {
 
 	private ISearchQueryExecutor doSearch(String theResourceType, SearchParameterMap theParams, ResourcePersistentId theReferencingPid) {
 		// keep this in sync with supportsSomeOf();
-		SearchSession session = getSearchSession();
+		if (theParams.getOffset() != null && theParams.getOffset() != 0) {
+			// perform an offset search instead of a scroll one, which doesn't allow for offset
+			List<Long> queryFetchResult = getSearchQueryOptionsStep(
+				theResourceType, theParams, theReferencingPid).fetchHits(theParams.getOffset(), theParams.getCount());
+			// indicate param was already processed, otherwise queries DB to process it
+			theParams.setOffset(null);
+			return SearchQueryExecutors.from(queryFetchResult);
+		}
 
+		SearchScroll<Long> esResult = getSearchScroll(theResourceType, theParams, theReferencingPid);
+		return new SearchScrollQueryExecutorAdaptor(esResult);
+	}
+
+
+	private SearchScroll<Long> getSearchScroll(String theResourceType, SearchParameterMap theParams, ResourcePersistentId theReferencingPid) {
 		int scrollSize = 50;
 		if (theParams.getCount()!=null) {
 			scrollSize = theParams.getCount();
 		}
 
-		SearchScroll<Long> esResult = session.search(ResourceTable.class)
+		return getSearchQueryOptionsStep(theResourceType, theParams, theReferencingPid).scroll(scrollSize);
+	}
+
+
+	private SearchQueryOptionsStep<?, Long, SearchLoadingOptionsStep, ?, ?> getSearchQueryOptionsStep(
+			String theResourceType, SearchParameterMap theParams, ResourcePersistentId theReferencingPid) {
+
+		return getSearchSession().search(ResourceTable.class)
 			// The document id is the PK which is pid.  We use this instead of _myId to avoid fetching the doc body.
 			.select(
 				// adapt the String docRef.id() to the Long that it really is.
@@ -188,11 +212,11 @@ public class FulltextSearchSvcImpl implements IFulltextSearchSvc {
 					//DROP EARLY HERE IF BOOL IS EMPTY?
 
 				})
-			).scroll(scrollSize);
+			);
-
-		return new SearchScrollQueryExecutorAdaptor(esResult);
 	}
 
+
+
 	@Nonnull
 	private SearchSession getSearchSession() {
 		return Search.session(myEntityManager);
@@ -314,4 +338,14 @@ public class FulltextSearchSvcImpl implements IFulltextSearchSvc {
 			.map(p -> p.toResource(parser))
 			.collect(Collectors.toList());
 	}
+
+
+
+	@Override
+	public long count(String theResourceName, SearchParameterMap theParams) {
+		SearchQueryOptionsStep<?, Long, SearchLoadingOptionsStep, ?, ?> queryOptionsStep =
+			getSearchQueryOptionsStep(theResourceName, theParams, null);
+
+		return queryOptionsStep.fetchTotalHitCount();
+	}
 }

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/dao/IFulltextSearchSvc.java

@@ -45,6 +45,7 @@ public interface IFulltextSearchSvc {
 	 */
 	List<ResourcePersistentId> search(String theResourceName, SearchParameterMap theParams);
 
+
 	/**
 	 * Query the index for a scrollable iterator of results.
 	 * No max size to the result iterator.
@@ -90,4 +91,8 @@ public interface IFulltextSearchSvc {
 	 */
 	List<IBaseResource> getResources(Collection<Long> thePids);
 
+	/**
+	 * Returns accurate hit count
+	 */
+	long count(String theResourceName, SearchParameterMap theParams);
 }

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/dao/ISearchBuilder.java

@@ -41,7 +41,7 @@ public interface ISearchBuilder {
 
 	IResultIterator createQuery(SearchParameterMap theParams, SearchRuntimeDetails theSearchRuntime, RequestDetails theRequest, @Nonnull RequestPartitionId theRequestPartitionId);
 
-	Iterator<Long> createCountQuery(SearchParameterMap theParams, String theSearchUuid, RequestDetails theRequest, RequestPartitionId theRequestPartitionId);
+	Long createCountQuery(SearchParameterMap theParams, String theSearchUuid, RequestDetails theRequest, RequestPartitionId theRequestPartitionId);
 
 	void setMaxResultsToFetch(Integer theMaxResultsToFetch);
 

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/dao/LegacySearchBuilder.java

@@ -227,14 +227,14 @@ public class LegacySearchBuilder implements ISearchBuilder {
 	}
 
 	@Override
-	public Iterator<Long> createCountQuery(SearchParameterMap theParams, String theSearchUuid, RequestDetails theRequest, @Nonnull RequestPartitionId theRequestPartitionId) {
+	public Long createCountQuery(SearchParameterMap theParams, String theSearchUuid, RequestDetails theRequest, @Nonnull RequestPartitionId theRequestPartitionId) {
 		assert theRequestPartitionId != null;
 		assert TransactionSynchronizationManager.isActualTransactionActive();
 
 		init(theParams, theSearchUuid, theRequestPartitionId);
 
 		List<TypedQuery<Long>> queries = createQuery(null, null, null, true, theRequest, null);
-		return new CountQueryIterator(queries.get(0));
+		return new CountQueryIterator(queries.get(0)).next();
 	}
 
 	/**

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/SearchCoordinatorSvcImpl.java

@@ -508,6 +508,7 @@ public class SearchCoordinatorSvcImpl implements ISearchCoordinatorSvc {
 		return candidate.orElse(null);
 	}
 
+
 	private IBundleProvider executeQuery(String theResourceType, SearchParameterMap theParams, RequestDetails theRequestDetails, String theSearchUuid, ISearchBuilder theSb, Integer theLoadSynchronousUpTo, RequestPartitionId theRequestPartitionId) {
 		SearchRuntimeDetails searchRuntimeDetails = new SearchRuntimeDetails(theRequestDetails, theSearchUuid);
 		searchRuntimeDetails.setLoadSynchronous(true);
@@ -533,12 +534,11 @@ public class SearchCoordinatorSvcImpl implements ISearchCoordinatorSvc {
 				List<List<IQueryParameterType>> contentAndTerms = theParams.get(Constants.PARAM_CONTENT);
 				List<List<IQueryParameterType>> textAndTerms = theParams.get(Constants.PARAM_TEXT);
 
-				Iterator<Long> countIterator = theSb.createCountQuery(theParams, theSearchUuid, theRequestDetails, theRequestPartitionId);
+				count = theSb.createCountQuery(theParams, theSearchUuid, theRequestDetails, theRequestPartitionId);
 
 				if (contentAndTerms != null) theParams.put(Constants.PARAM_CONTENT, contentAndTerms);
 				if (textAndTerms != null) theParams.put(Constants.PARAM_TEXT, textAndTerms);
 
-				count = countIterator.next();
 				ourLog.trace("Got count {}", count);
 			}
 
@@ -1233,8 +1233,8 @@ public class SearchCoordinatorSvcImpl implements ISearchCoordinatorSvc {
 				 * we will have to clone those parameters here so that
 				 * the "correct" params are used in createQuery below
 				 */
-				Iterator<Long> countIterator = sb.createCountQuery(myParams.clone(), mySearch.getUuid(), myRequest, myRequestPartitionId);
+				Long count = sb.createCountQuery(myParams.clone(), mySearch.getUuid(), myRequest, myRequestPartitionId);
-				Long count = countIterator.hasNext() ? countIterator.next() : 0L;
+
 				ourLog.trace("Got count {}", count);
 
 				TransactionTemplate txTemplate = new TransactionTemplate(myManagedTxManager);

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/builder/SearchBuilder.java

@@ -75,6 +75,7 @@ import ca.uhn.fhir.model.valueset.BundleEntrySearchModeEnum;
 import ca.uhn.fhir.rest.api.Constants;
 import ca.uhn.fhir.rest.api.RestSearchParameterTypeEnum;
 import ca.uhn.fhir.rest.api.SearchContainedModeEnum;
+import ca.uhn.fhir.rest.api.SearchTotalModeEnum;
 import ca.uhn.fhir.rest.api.SortOrderEnum;
 import ca.uhn.fhir.rest.api.SortSpec;
 import ca.uhn.fhir.rest.api.server.IPreResourceAccessDetails;
@@ -97,6 +98,7 @@ import com.google.common.collect.Streams;
 import com.healthmarketscience.sqlbuilder.Condition;
 import org.apache.commons.lang3.Validate;
 import org.apache.commons.lang3.math.NumberUtils;
+import org.apache.jena.sparql.engine.QueryIterator;
 import org.hl7.fhir.instance.model.api.IAnyResource;
 import org.hl7.fhir.instance.model.api.IBaseResource;
 import org.slf4j.Logger;
@@ -119,6 +121,7 @@ import javax.persistence.criteria.From;
 import javax.persistence.criteria.Predicate;
 import javax.persistence.criteria.Root;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -271,18 +274,24 @@ public class SearchBuilder implements ISearchBuilder {
 
 	@SuppressWarnings("ConstantConditions")
 	@Override
-	public Iterator<Long> createCountQuery(SearchParameterMap theParams, String theSearchUuid, RequestDetails theRequest, @Nonnull RequestPartitionId theRequestPartitionId) {
+	public Long createCountQuery(SearchParameterMap theParams, String theSearchUuid,
+				RequestDetails theRequest, @Nonnull RequestPartitionId theRequestPartitionId) {
+
 		assert theRequestPartitionId != null;
 		assert TransactionSynchronizationManager.isActualTransactionActive();
 
 		init(theParams, theSearchUuid, theRequestPartitionId);
 
-		List<ISearchQueryExecutor> queries = createQuery(myParams, null, null, null, true, theRequest, null);
+		if (checkUseHibernateSearch()) {
-		if (queries.isEmpty()) {
+			long count = myFulltextSearchSvc.count(myResourceName, theParams.clone());
-			return Collections.emptyIterator();
+			return count;
 		}
-		try (ISearchQueryExecutor queryExecutor = queries.get(0)) {
+
-			return Lists.newArrayList(queryExecutor.next()).iterator();
+		List<ISearchQueryExecutor> queries = createQuery(theParams.clone(), null, null, null, true, theRequest, null);
+		if (queries.isEmpty()) {
+			return 0L;
+		} else {
+			return queries.get(0).next();
 		}
 	}
 
@@ -309,6 +318,7 @@ public class SearchBuilder implements ISearchBuilder {
 		return new QueryIterator(theSearchRuntimeDetails, theRequest);
 	}
 
+
 	private void init(SearchParameterMap theParams, String theSearchUuid, RequestPartitionId theRequestPartitionId) {
 		myCriteriaBuilder = myEntityManager.getCriteriaBuilder();
 		myParams = theParams;
@@ -316,7 +326,7 @@ public class SearchBuilder implements ISearchBuilder {
 		myRequestPartitionId = theRequestPartitionId;
 	}
 
-	private List<ISearchQueryExecutor> createQuery(SearchParameterMap theParams, SortSpec sort, Integer theOffset, Integer theMaximumResults, boolean theCount, RequestDetails theRequest,
+	private List<ISearchQueryExecutor> createQuery(SearchParameterMap theParams, SortSpec sort, Integer theOffset, Integer theMaximumResults, boolean theCountOnlyFlag, RequestDetails theRequest,
 																		 SearchRuntimeDetails theSearchRuntimeDetails) {
 
 		ArrayList<ISearchQueryExecutor> queries = new ArrayList<>();
@@ -359,8 +369,6 @@ public class SearchBuilder implements ISearchBuilder {
 				!fulltextExecutor.hasNext() ||
 					// Our hibernate search query doesn't respect partitions yet
 					(!myPartitionSettings.isPartitioningEnabled() &&
-						// we don't support _count=0 yet.
-						!theCount &&
 					// were there AND terms left?  Then we still need the db.
 						theParams.isEmpty() &&
 						// not every param is a param. :-(
@@ -382,11 +390,11 @@ public class SearchBuilder implements ISearchBuilder {
 				// We break the pids into chunks that fit in the 1k limit for jdbc bind params.
 				// wipmb change chunk to take iterator
 				new QueryChunker<Long>()
-					.chunk(Streams.stream(fulltextExecutor).collect(Collectors.toList()), t -> doCreateChunkedQueries(theParams, t, theOffset, sort, theCount, theRequest, queries));
+					.chunk(Streams.stream(fulltextExecutor).collect(Collectors.toList()), t -> doCreateChunkedQueries(theParams, t, theOffset, sort, theCountOnlyFlag, theRequest, queries));
 			}
 		} else {
 			// do everything in the database.
-			Optional<SearchQueryExecutor> query = createChunkedQuery(theParams, sort, theOffset, theMaximumResults, theCount, theRequest, null);
+			Optional<SearchQueryExecutor> query = createChunkedQuery(theParams, sort, theOffset, theMaximumResults, theCountOnlyFlag, theRequest, null);
 			query.ifPresent(queries::add);
 		}
 
@@ -408,7 +416,7 @@ public class SearchBuilder implements ISearchBuilder {
 		}
 
 		// TODO MB someday we'll want a query planner to figure out if we _should_ or _must_ use the ft index, not just if we can.
-		return fulltextEnabled &&
+		return fulltextEnabled && myParams != null &&
 			myParams.getSearchContainedMode() == SearchContainedModeEnum.FALSE &&
 			myFulltextSearchSvc.supportsSomeOf(myParams);
 	}
@@ -506,9 +514,9 @@ public class SearchBuilder implements ISearchBuilder {
 		}
 	}
 
-	private Optional<SearchQueryExecutor> createChunkedQuery(SearchParameterMap theParams, SortSpec sort, Integer theOffset, Integer theMaximumResults, boolean theCount, RequestDetails theRequest, List<Long> thePidList) {
+	private Optional<SearchQueryExecutor> createChunkedQuery(SearchParameterMap theParams, SortSpec sort, Integer theOffset, Integer theMaximumResults, boolean theCountOnlyFlag, RequestDetails theRequest, List<Long> thePidList) {
 		String sqlBuilderResourceName = myParams.getEverythingMode() == null ? myResourceName : null;
-		SearchQueryBuilder sqlBuilder = new SearchQueryBuilder(myContext, myDaoConfig.getModelConfig(), myPartitionSettings, myRequestPartitionId, sqlBuilderResourceName, mySqlBuilderFactory, myDialectProvider, theCount);
+		SearchQueryBuilder sqlBuilder = new SearchQueryBuilder(myContext, myDaoConfig.getModelConfig(), myPartitionSettings, myRequestPartitionId, sqlBuilderResourceName, mySqlBuilderFactory, myDialectProvider, theCountOnlyFlag);
 		QueryStack queryStack3 = new QueryStack(theParams, myDaoConfig, myDaoConfig.getModelConfig(), myContext, sqlBuilder, mySearchParamRegistry, myPartitionSettings);
 
 		if (theParams.keySet().size() > 1 || theParams.getSort() != null || theParams.keySet().contains(Constants.PARAM_HAS) || isPotentiallyContainedReferenceParameterExistsAtRoot(theParams)) {
@@ -533,7 +541,7 @@ public class SearchBuilder implements ISearchBuilder {
 				// is basically a reverse-include search. For type/Everything (as opposed to instance/Everything)
 				// the one problem with this approach is that it doesn't catch Patients that have absolutely
 				// nothing linked to them. So we do one additional query to make sure we catch those too.
-				SearchQueryBuilder fetchPidsSqlBuilder = new SearchQueryBuilder(myContext, myDaoConfig.getModelConfig(), myPartitionSettings, myRequestPartitionId, myResourceName, mySqlBuilderFactory, myDialectProvider, theCount);
+				SearchQueryBuilder fetchPidsSqlBuilder = new SearchQueryBuilder(myContext, myDaoConfig.getModelConfig(), myPartitionSettings, myRequestPartitionId, myResourceName, mySqlBuilderFactory, myDialectProvider, theCountOnlyFlag);
 				GeneratedSql allTargetsSql = fetchPidsSqlBuilder.generate(theOffset, myMaxResultsToFetch);
 				String sql = allTargetsSql.getSql();
 				Object[] args = allTargetsSql.getBindVariables().toArray(new Object[0]);
@@ -613,7 +621,7 @@ public class SearchBuilder implements ISearchBuilder {
 		 * finds the appropriate resources) in an outer search which is then sorted
 		 */
 		if (sort != null) {
-			assert !theCount;
+			assert !theCountOnlyFlag;
 
 			createSort(queryStack3, sort);
 		}

hapi-fhir-jpaserver-searchparam/src/main/java/ca/uhn/fhir/jpa/searchparam/MatchUrlService.java

@@ -33,6 +33,7 @@ import ca.uhn.fhir.model.api.Include;
 import ca.uhn.fhir.rest.api.Constants;
 import ca.uhn.fhir.rest.api.QualifiedParamList;
 import ca.uhn.fhir.rest.api.RestSearchParameterTypeEnum;
+import ca.uhn.fhir.rest.api.SearchTotalModeEnum;
 import ca.uhn.fhir.rest.param.DateRangeParam;
 import ca.uhn.fhir.rest.param.ParameterUtil;
 import ca.uhn.fhir.rest.server.exceptions.InvalidRequestException;
@@ -113,6 +114,15 @@ public class MatchUrlService {
 						throw new InvalidRequestException(Msg.code(485) + "Invalid " + Constants.PARAM_COUNT + " value: " + intString);
 					}
 				}
+			} else if (Constants.PARAM_SEARCH_TOTAL_MODE.equals(nextParamName)) {
+				if (paramList != null && ! paramList.isEmpty() && ! paramList.get(0).isEmpty()) {
+					String totalModeEnumStr = paramList.get(0).get(0);
+					try {
+						paramMap.setSearchTotalMode(SearchTotalModeEnum.valueOf(totalModeEnumStr));
+					} catch (IllegalArgumentException e) {
+						throw new InvalidRequestException(Msg.code(2078) + "Invalid " + Constants.PARAM_SEARCH_TOTAL_MODE + " value: " + totalModeEnumStr);
+					}
+				}
 			} else if (Constants.PARAM_OFFSET.equals(nextParamName)) {
 				if (paramList != null && paramList.size() > 0 && paramList.get(0).size() > 0) {
 					String intString = paramList.get(0).get(0);

hapi-fhir-jpaserver-test-utilities/src/test/java/ca/uhn/fhir/jpa/dao/TestDaoSearch.java

@@ -19,6 +19,7 @@ import org.springframework.web.util.UriComponentsBuilder;
 
 import javax.annotation.Nonnull;
 import java.util.List;
+import java.util.stream.Collectors;
 
 /**
  * Simplistic implementation of FHIR queries.
@@ -51,11 +52,13 @@ public class TestDaoSearch {
 		return result.getAllResources();
 	}
 
-	public List<String> searchForIds(String theQueryUrl) {
+	public List<String>  searchForIds(String theQueryUrl) {
 		// fake out the server url parsing
 		IBundleProvider result = searchForBundleProvider(theQueryUrl);
 
-		List<String> resourceIds = result.getAllResourceIds();
+		// getAllResources is not safe as size is not always set
+		List<String> resourceIds = result.getResources(0, Integer.MAX_VALUE)
+					.stream().map(resource -> resource.getIdElement().getIdPart()).collect(Collectors.toList());
 		return resourceIds;
 	}
 

hapi-fhir-jpaserver-test-utilities/src/test/java/ca/uhn/fhir/jpa/dao/r4/FhirResourceDaoR4SearchWithElasticSearchIT.java

@@ -30,6 +30,7 @@ import ca.uhn.fhir.jpa.test.config.TestR4Config;
 import ca.uhn.fhir.parser.DataFormatException;
 import ca.uhn.fhir.parser.IParser;
 import ca.uhn.fhir.rest.api.Constants;
+import ca.uhn.fhir.rest.api.SearchTotalModeEnum;
 import ca.uhn.fhir.rest.api.server.IBundleProvider;
 import ca.uhn.fhir.rest.api.server.storage.ResourcePersistentId;
 import ca.uhn.fhir.rest.param.ReferenceParam;
@@ -45,6 +46,7 @@ import ca.uhn.fhir.test.utilities.LogbackLevelOverrideExtension;
 import ca.uhn.fhir.test.utilities.docker.RequiresDocker;
 import ca.uhn.fhir.validation.FhirValidator;
 import ca.uhn.fhir.validation.ValidationResult;
+import com.google.common.collect.Lists;
 import org.hamcrest.Matchers;
 import org.hl7.fhir.instance.model.api.IBaseCoding;
 import org.hl7.fhir.instance.model.api.IBaseResource;
@@ -71,6 +73,8 @@ import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.test.annotation.DirtiesContext;
@@ -84,8 +88,11 @@ import org.springframework.transaction.PlatformTransactionManager;
 
 import javax.persistence.EntityManager;
 import java.io.IOException;
+import java.time.Month;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 import static ca.uhn.fhir.jpa.model.util.UcumServiceUtil.UCUM_CODESYSTEM_URL;
@@ -99,6 +106,7 @@ import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.stringContainsInOrder;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -1472,9 +1480,9 @@ public class FhirResourceDaoR4SearchWithElasticSearchIT extends BaseJpaTest {
 		}
 
 		private IIdType withObservationWithQuantity(double theValue, String theSystem, String theCode) {
-			myResourceId = myTestDataBuilder.createObservation(
+			myResourceId = myTestDataBuilder.createObservation(asArray(
 				myTestDataBuilder.withQuantityAtPath("valueQuantity", theValue, theSystem, theCode)
-			);
+			));
 			return myResourceId;
 		}
 
@@ -1490,6 +1498,87 @@ public class FhirResourceDaoR4SearchWithElasticSearchIT extends BaseJpaTest {
 	}
 
 
+
+	@Nested
+	public class TotalParameter {
+
+		@ParameterizedTest
+		@EnumSource(SearchTotalModeEnum.class)
+		public void totalParamSkipsSql(SearchTotalModeEnum theTotalModeEnum) {
+			myTestDataBuilder.createObservation(asArray(myTestDataBuilder.withObservationCode("http://example.com/", "theCode")));
+
+			myCaptureQueriesListener.clear();
+			myTestDaoSearch.searchForIds("Observation?code=theCode&_total=" + theTotalModeEnum);
+			myCaptureQueriesListener.logSelectQueriesForCurrentThread();
+			assertEquals(1, myCaptureQueriesListener.getSelectQueriesForCurrentThread().size(), "bundle was built with no sql");
+		}
+
+
+		@Test
+		public void totalIsCorrect() {
+			myTestDataBuilder.createObservation(asArray(myTestDataBuilder.withObservationCode("http://example.com/", "code-1")));
+			myTestDataBuilder.createObservation(asArray(myTestDataBuilder.withObservationCode("http://example.com/", "code-2")));
+			myTestDataBuilder.createObservation(asArray(myTestDataBuilder.withObservationCode("http://example.com/", "code-3")));
+
+			IBundleProvider resultBundle = myTestDaoSearch.searchForBundleProvider("Observation?_total=" + SearchTotalModeEnum.ACCURATE);
+			assertEquals(3, resultBundle.size());
+		}
+
+	}
+
+
+	@Nested
+	public class OffsetParameter {
+
+		@BeforeEach
+		public void enableResourceStorage() {
+			myDaoConfig.setStoreResourceInLuceneIndex(true);
+		}
+
+
+		@Test
+		public void offsetNoCount() {
+			myTestDataBuilder.createObservation(asArray(myTestDataBuilder.withObservationCode("http://example.com/", "code-1")));
+			IIdType idCode2 = myTestDataBuilder.createObservation(asArray(myTestDataBuilder.withObservationCode("http://example.com/", "code-2")));
+			IIdType idCode3 = myTestDataBuilder.createObservation(asArray(myTestDataBuilder.withObservationCode("http://example.com/", "code-3")));
+
+			myCaptureQueriesListener.clear();
+			List<String> resultIds = myTestDaoSearch.searchForIds("Observation?code=code-1,code-2,code-3&_offset=1");
+			myCaptureQueriesListener.logSelectQueriesForCurrentThread();
+
+			assertThat(resultIds, containsInAnyOrder(idCode2.getIdPart(), idCode3.getIdPart()));
+			// make also sure no extra SQL queries were executed
+			assertEquals(0, myCaptureQueriesListener.getSelectQueriesForCurrentThread().size(), "bundle was built with no sql");
+		}
+
+
+		@Test
+		public void offsetAndCount() {
+			myTestDataBuilder.createObservation(asArray(myTestDataBuilder.withObservationCode("http://example.com/", "code-1")));
+			IIdType idCode2 = myTestDataBuilder.createObservation(asArray(myTestDataBuilder.withObservationCode("http://example.com/", "code-2")));
+			myTestDataBuilder.createObservation(asArray(myTestDataBuilder.withObservationCode("http://example.com/", "code-3")));
+
+			myCaptureQueriesListener.clear();
+			List<String> resultIds = myTestDaoSearch.searchForIds("Observation?code=code-1,code-2,code-3&_offset=1&_count=1");
+			myCaptureQueriesListener.logSelectQueriesForCurrentThread();
+
+			assertThat(resultIds, containsInAnyOrder(idCode2.getIdPart()));
+			// also validate no extra SQL queries were executed
+			assertEquals(0, myCaptureQueriesListener.getSelectQueriesForCurrentThread().size(), "bundle was built with no sql");
+		}
+	}
+
+
+	private Consumer<IBaseResource>[] asArray(Consumer<IBaseResource> theIBaseResourceConsumer) {
+		@SuppressWarnings("unchecked")
+		Consumer<IBaseResource>[] array = (Consumer<IBaseResource>[]) new Consumer[]{theIBaseResourceConsumer};
+		return array;
+	}
+
+
+
+
+
 	/**
 	 * Disallow context dirtying for nested classes
 	 */

hapi-fhir-jpaserver-test-utilities/src/test/java/ca/uhn/fhir/jpa/provider/r4/ResourceProviderR4ElasticTest.java

@@ -36,6 +36,7 @@ import java.time.Instant;
 import java.util.Date;
 import java.util.List;
 import java.util.Objects;
+import java.util.stream.IntStream;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.hasItem;
@@ -43,6 +44,7 @@ import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.nullValue;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 @ExtendWith(SpringExtension.class)
@@ -148,7 +150,7 @@ public class ResourceProviderR4ElasticTest extends BaseResourceProviderR4Test {
 			.useHttpGet()
 			.execute();
 
-		assertEquals( 1, respParam.getParameter().size(), "Expected only 1 observation for blood count code");
+		assertEquals(1, respParam.getParameter().size(), "Expected only 1 observation for blood count code");
 		Bundle bundle = (Bundle) respParam.getParameter().get(0).getResource();
 		Observation observation = (Observation) bundle.getEntryFirstRep().getResource();
 
@@ -157,4 +159,47 @@ public class ResourceProviderR4ElasticTest extends BaseResourceProviderR4Test {
 
 	}
 
+	@Test
+	public void testCountReturnsExpectedSizeOfResources() throws IOException {
+		IntStream.range(0, 10).forEach(index -> {
+			Coding blood_count = new Coding("http://loinc.org", "789-8", "Erythrocytes in Blood by Automated count for code: " + (index + 1));
+			createObservationWithCode(blood_count);
+		});
+		HttpGet countQuery = new HttpGet(ourServerBase + "/Observation?code=789-8&_count=5");
+		myCaptureQueriesListener.clear();
+		try (CloseableHttpResponse response = ourHttpClient.execute(countQuery)) {
+			myCaptureQueriesListener.logSelectQueriesForCurrentThread();
+			// then
+			assertEquals(Constants.STATUS_HTTP_200_OK, response.getStatusLine().getStatusCode());
+			String text = IOUtils.toString(response.getEntity().getContent(), StandardCharsets.UTF_8);
+			Bundle bundle = myFhirContext.newXmlParser().parseResource(Bundle.class, text);
+			assertEquals(10, bundle.getTotal(), "Expected total 10 observations matching query");
+			assertEquals(5, bundle.getEntry().size(), "Expected 5 observation entries to match page size");
+			assertTrue(bundle.getLink("next").hasRelation());
+			assertEquals(0, myCaptureQueriesListener.getSelectQueriesForCurrentThread().size(), "we build the bundle with no sql");
+		}
+	}
+
+	@Test
+	public void testCountZeroReturnsNoResourceEntries() throws IOException {
+		IntStream.range(0, 10).forEach(index -> {
+			Coding blood_count = new Coding("http://loinc.org", "789-8", "Erythrocytes in Blood by Automated count for code: " + (index + 1));
+			createObservationWithCode(blood_count);
+		});
+		HttpGet countQuery = new HttpGet(ourServerBase + "/Observation?code=789-8&_count=0");
+		myCaptureQueriesListener.clear();
+		try (CloseableHttpResponse response = ourHttpClient.execute(countQuery)) {
+			myCaptureQueriesListener.logSelectQueriesForCurrentThread();
+			assertEquals(Constants.STATUS_HTTP_200_OK, response.getStatusLine().getStatusCode());
+			String text = IOUtils.toString(response.getEntity().getContent(), StandardCharsets.UTF_8);
+			Bundle bundle = myFhirContext.newXmlParser().parseResource(Bundle.class, text);
+			assertEquals(10, bundle.getTotal(), "Expected total 10 observations matching query");
+			assertEquals(0, bundle.getEntry().size(), "Expected no entries in bundle");
+			assertNull(bundle.getLink("next"), "Expected no 'next' link");
+			assertNull(bundle.getLink("prev"), "Expected no 'prev' link");
+			assertEquals(0, myCaptureQueriesListener.getSelectQueriesForCurrentThread().size(), "we build the bundle with no sql");
+		}
+
+	}
+
 }

hapi-fhir-jpaserver-test-utilities/src/test/java/ca/uhn/fhir/jpa/search/SearchCoordinatorSvcImplTest.java

@@ -561,7 +561,7 @@ public class SearchCoordinatorSvcImplTest {
 		params.setSearchTotalMode(SearchTotalModeEnum.ACCURATE);
 
 		List<ResourcePersistentId> pids = createPidSequence(30);
-		when(mySearchBuilder.createCountQuery(same(params), any(String.class), nullable(RequestDetails.class), nullable(RequestPartitionId.class))).thenReturn(Lists.newArrayList(Long.valueOf(20L)).iterator());
+		when(mySearchBuilder.createCountQuery(same(params), any(String.class),nullable(RequestDetails.class), nullable(RequestPartitionId.class))).thenReturn(20L);
 		when(mySearchBuilder.createQuery(same(params), any(), nullable(RequestDetails.class), nullable(RequestPartitionId.class))).thenReturn(new ResultIterator(pids.subList(10, 20).iterator()));
 
 		doAnswer(loadPids()).when(mySearchBuilder).loadResourcesByPid(any(Collection.class), any(Collection.class), any(List.class), anyBoolean(), any());