Replace page limit with entry limit

cgpclient/client.py

@@ -636,10 +636,11 @@ def get_referrals(self) -> CGPReferrals:
             client=self,
         )
 
-    def get_files(self) -> CGPFiles:
+    def get_files(self, max_results=100) -> CGPFiles:
         return CGPFiles(
             document_references=self.fhir_service.search_for_document_references(
-                search_params=self.fhir_config
+                search_params=self.fhir_config,
+                max_results=max_results
             ),
             client=self,
         )

cgpclient/fhir.py

@@ -50,11 +50,10 @@
 
 log = logging.getLogger(__name__)
 
-MAX_SEARCH_RESULTS = 100
+DEFAULT_MAX_SEARCH_RESULTS = 100
 MAX_UNSIGNED_INT = (
     2147483647  # https://hl7.org/fhir/R4/datatypes.html#unsignedInt # noqa: E501
 )
-MAX_PAGES = 100
 
 
 # Enumerations for various FHIR resource fields
@@ -196,10 +195,12 @@ def _search_paged(
         self,
         url: str,
         query_params: list[tuple] | None = None,
+        max_results: int = DEFAULT_MAX_SEARCH_RESULTS
     ) -> Bundle:
         """Peform a search request and page through the results"""
         pages = 1
-        while pages <= MAX_PAGES:
+        entries = 0
+        while entries <= max_results:
             response = requests.get(
                 url=url,
                 headers=self.headers,
@@ -208,13 +209,14 @@ def _search_paged(
             )
             if response.ok:
                 bundle = Bundle.parse_obj(response.json())
+                entries += len(bundle.entry) if bundle.entry else 0
                 yield bundle
                 if (
                     bundle.link
                     and len(bundle.link) == 1
                     and bundle.link[0].relation == "next"
                 ):
-                    log.info("Requesting page %i", pages)
+                    log.info("Requesting page %i, entries so far %i", pages, entries)
                     # we need to switch the HealthLake URL for one with our prefix # noqa: E501
                     prefix = url.rsplit("/", 1)[0]
                     suffix = bundle.link[0].url.rsplit("/", 1)[-1]
@@ -225,7 +227,7 @@ def _search_paged(
                     return
             else:
                 raise CGPClientException(f"Failed to fetch from endpoint: {url}")
-        log.info("Reached maximum number of pages")
+        raise CGPClientException("Reached maximum number of entries")
 
     def _merge_bundles(self, bundles: list[Bundle]) -> Bundle:
         """Merge a list of Bundles into a single one, retaining the
@@ -236,16 +238,14 @@ def _merge_bundles(self, bundles: list[Bundle]) -> Bundle:
         return first
 
     def search_for_fhir_resource(
-        self, resource_type: str, query_params: list[tuple] = []
+        self, resource_type: str, query_params: list[tuple] = [], max_results: int = DEFAULT_MAX_SEARCH_RESULTS
     ) -> Bundle:
         """Search for a FHIR resource using the query parameters"""
         url = f"{self.base_url}/{resource_type}"
 
         if query_params is None:
             query_params = []
 
-        query_params.append(("_count", str(MAX_SEARCH_RESULTS)))
-
         if self.config.workspace_id is not None:
             query_params.append(("_tag", self.config.workspace_id))
 
@@ -254,7 +254,7 @@ def search_for_fhir_resource(
 
         bundles: list[Bundle] = []
 
-        for response in self._search_paged(url=url, query_params=query_params):
+        for response in self._search_paged(url=url, query_params=query_params, max_results=max_results):
             bundles.append(response)
 
         return self._merge_bundles(bundles)
@@ -283,7 +283,7 @@ def search_for_tasks(self, search_params: FHIRConfig | None = None) -> list[Task
                 )
 
     def search_for_document_references(
-        self, search_params: FHIRConfig | None = None
+        self, search_params: FHIRConfig | None = None, max_results: int = DEFAULT_MAX_SEARCH_RESULTS
     ) -> list[DocumentReference]:
         """Search for DocumentReferences using the parameters in the FHIR
         config"""
@@ -333,6 +333,7 @@ def search_for_document_references(
         bundle: Bundle = self.search_for_fhir_resource(
             resource_type=DocumentReference.__name__,
             query_params=query_params,
+            max_results=max_results
         )
 
         if bundle.entry:

tests/test_fhir.py

@@ -8,6 +8,7 @@
 
 from cgpclient.client import CGPClient
 from cgpclient.fhir import CGPFHIRClient, FHIRConfig  # type: ignore
+from cgpclient.utils import CGPClientException
 
 
 @pytest.fixture(scope="function")
@@ -261,6 +262,27 @@ def json(self):
     assert resource.entry and len(resource.entry) == 1
 
 
+@patch("cgpclient.fhir.requests.get")
+def test_search_limits_resutls(mock_get: MagicMock, doc_ref_bundle: dict) -> None:
+    class MockedResponse:
+        def ok(self):
+            return True
+
+        def json(self):
+            doc_ref_bundle['link'] = [{'relation': 'next', 'url': '/'}]
+            return doc_ref_bundle
+
+    mock_get.side_effect = [MockedResponse()] * 10
+
+    config: FHIRConfig = FHIRConfig()
+
+    fhir: CGPFHIRClient = CGPFHIRClient(
+        api_base_url="host", headers={}, config=config, dry_run=False
+    )
+    with pytest.raises(CGPClientException):
+        fhir.search_for_fhir_resource(resource_type="DocumentReference", max_results=5)
+
+
 @patch("cgpclient.fhir.requests.get")
 def test_search_doc_refs(mock_get: MagicMock, doc_ref_bundle: dict) -> None:
     class MockedResponse: