Expand infinite loop fix to have hardcoded guard rail for max requests (#2250)

* Expand infinite loop fix to have hardcoded guard rail

* Add simple tests

* Fix docstring typo

Author: ml-evs

optimade/client/client.py

@@ -90,6 +90,12 @@ class OptimadeClient:
     max_attempts: int
     """The maximum number of times to repeat a failed query before giving up."""
 
+    max_requests_per_provider: int = 2_000_000
+    """An upper limit guard rail to avoid infinite hanging of the client on malformed APIs.
+    If available, a better value will be estimated for each API based on the total number of entries.
+
+    """
+
     use_async: bool
     """Whether or not to make all requests asynchronously using asyncio."""
 
@@ -957,9 +963,13 @@ async def _get_one_async(
         request_delay: float | None = None
 
         results = QueryResults()
+        number_of_requests = 0
+        total_data_available: int | None = None
         try:
             async with self._http_client(headers=self.headers) as client:  # type: ignore[union-attr,call-arg,misc]
                 while next_url:
+                    number_of_requests += 1
+
                     attempts = 0
                     try:
                         if self.verbosity:
@@ -975,6 +985,21 @@ async def _get_one_async(
                         if request_delay:
                             request_delay = min(request_delay, 5)
 
+                        # Compute the upper limit guard rail on pagination requests based on the number of entries in the entire db
+                        # and the chosen page limit
+                        if total_data_available is None:
+                            total_data_available = page_results["meta"].get(
+                                "data_available", 0
+                            )
+                            page_limit = len(page_results["data"])
+                            if total_data_available and total_data_available > 0:
+                                stopping_criteria = min(
+                                    math.ceil(total_data_available / page_limit),
+                                    self.max_requests_per_provider,
+                                )
+                            else:
+                                stopping_criteria = self.max_results_per_provider
+
                     except RecoverableHTTPError:
                         attempts += 1
                         if attempts > self.max_attempts:
@@ -989,9 +1014,9 @@ async def _get_one_async(
                     if not paginate:
                         break
 
-                    if len(results.data) == 0:
+                    if len(results.data) == 0 or number_of_requests > stopping_criteria:
                         if next_url:
-                            message = f"{base_url} unexpectedly stopped returning results. Stopping download."
+                            message = f"Detected potential infinite loop for {base_url} (more than {stopping_criteria=} requests made). Stopping download."
                             results.errors.append(message)
                             if not self.silent:
                                 self._progress.print(message)
@@ -1041,6 +1066,8 @@ def _get_one(
         request_delay: float | None = None
 
         results = QueryResults()
+        number_of_requests: int = 0
+        total_data_available: int | None = None
         try:
             with self._http_client() as client:  # type: ignore[misc]
                 client.headers.update(self.headers)
@@ -1050,6 +1077,7 @@ def _get_one(
                     timeout = (self.http_timeout.connect, self.http_timeout.read)
 
                 while next_url:
+                    number_of_requests += 1
                     attempts = 0
                     try:
                         if self.verbosity:
@@ -1059,6 +1087,21 @@ def _get_one(
                         r = client.get(next_url, timeout=timeout)
                         page_results, next_url = self._handle_response(r, _task)
 
+                        # Compute the upper limit guard rail on pagination requests based on the number of entries in the entire db
+                        # and the chosen page limit
+                        if total_data_available is None:
+                            total_data_available = page_results["meta"].get(
+                                "data_available", 0
+                            )
+                            page_limit = len(page_results["data"])
+                            if total_data_available and total_data_available > 0:
+                                stopping_criteria = min(
+                                    math.ceil(total_data_available / page_limit),
+                                    self.max_requests_per_provider,
+                                )
+                            else:
+                                stopping_criteria = self.max_results_per_provider
+
                         request_delay = page_results["meta"].get("request_delay", None)
                         # Don't wait any longer than 5 seconds
                         if request_delay:
@@ -1075,9 +1118,9 @@ def _get_one(
 
                     results.update(page_results)
 
-                    if len(results.data) == 0:
+                    if len(results.data) == 0 or number_of_requests > stopping_criteria:
                         if next_url:
-                            message = f"{base_url} unexpectedly stopped returning results. Stopping download."
+                            message = f"Detected potential infinite loop for {base_url} (more than {stopping_criteria=} requests made). Stopping download."
                             results.errors.append(message)
                             if not self.silent:
                                 self._progress.print(message)

tests/server/test_client.py

@@ -599,3 +599,34 @@ async def test_raw_get_one_async(async_http_client):
         paginate=False,
     )
     assert len(override[TEST_URL].data) == 1
+
+
+@pytest.mark.asyncio
+async def test_guardrail_async(async_http_client):
+    """Test the upper limit on requests guard rail."""
+    cli = OptimadeClient(
+        base_urls=[TEST_URL],
+        http_client=async_http_client,
+        use_async=True,
+    )
+    cli.max_requests_per_provider = 1
+    result = await cli.get_one_async(
+        endpoint="structures", base_url=TEST_URL, filter="", page_limit=5, paginate=True
+    )
+    assert len(result[TEST_URL].errors) == 1
+    assert "infinite" in result[TEST_URL].errors[0]
+
+
+def test_guardrail_sync(http_client):
+    """Test the upper limit on requests guard rail."""
+    cli = OptimadeClient(
+        base_urls=[TEST_URL],
+        http_client=http_client,
+        use_async=False,
+    )
+    cli.max_requests_per_provider = 1
+    result = cli.get_one(
+        endpoint="structures", base_url=TEST_URL, filter="", page_limit=5, paginate=True
+    )
+    assert len(result[TEST_URL].errors) == 1
+    assert "infinite" in result[TEST_URL].errors[0]