fix bugs in batched retrieval

Author: esoteric-ephemera

mp_api/client/core/client.py

@@ -805,27 +805,29 @@ def _query_resource(
         except RequestException as ex:
             raise MPRestError(str(ex))
 
-    def _submit_requests(  # noqa
+    def _submit_requests(
         self,
-        url,
-        criteria,
-        use_document_model,
-        chunk_size,
-        num_chunks=None,
-        timeout=None,
+        url: str,
+        criteria: dict[str, Any],
+        use_document_model: bool,
+        chunk_size: int,
+        num_chunks: int | None = None,
+        timeout: float | None = None,
+        max_batch_size: int = 100,
     ) -> dict:
         """Handle submitting requests sequentially with pagination.
 
         If criteria contains comma-separated parameters (except those that are naturally comma-separated),
         split them into multiple sequential requests and combine results.
 
         Arguments:
-            criteria: dictionary of criteria to filter down
-            url: url used to make request
-            use_document_model: if None, will defer to the self.use_document_model attribute
-            num_chunks: Maximum number of chunks of data to yield. None will yield all possible.
-            chunk_size: Number of data entries per chunk.
-            timeout: Time in seconds to wait until a request timeout error is thrown
+            url (str): url used to make request
+            criteria (dict of str): dictionary of criteria to filter down
+            use_document_model (bool): whether to use the document model
+            num_chunks (int or None): Maximum number of chunks of data to yield. None will yield all possible.
+            chunk_size (int): Number of data entries per chunk.
+            timeout (float): Time in seconds to wait until a request timeout error is thrown
+            max_batch_size (int) : Maximum size of a batch for when retrieving batches in parallel
 
         Returns:
             Dictionary containing data and metadata
@@ -884,14 +886,6 @@ def _submit_requests(  # noqa
                     timeout=timeout,
                 )
 
-                # Check if we got 0 results - some parameters are silently ignored by the API
-                # when passed as comma-separated values, so we need to split them anyway
-                if total_num_docs == 0 and len(split_values) > 1:
-                    # Treat this the same as a 422 error - split into batches
-                    raise MPRestError(
-                        "Got 0 results for comma-separated parameter, will try splitting"
-                    )
-
                 # If successful, continue with normal pagination
                 data_chunks = [data["data"]]
                 total_data: dict[str, Any] = {"data": []}
@@ -904,17 +898,25 @@ def _submit_requests(  # noqa
 
             except MPRestError as e:
                 # If we get 422 or 414 error, or 0 results for comma-separated params, split into batches
-                if any(trace in str(e) for trace in ("422", "414", "Got 0 results")):
+                if any(
+                    trace in str(e)
+                    for trace in (
+                        "422",
+                        "414",
+                    )
+                ):
                     total_data = {"data": []}
                     total_num_docs = 0
                     data_chunks = []
 
                     # Batch the split values to reduce number of requests
                     # Use batches of up to 100 values to balance URL length and request count
-                    batch_size = min(100, max(1, len(split_values) // 10))
+                    num_batches = min(
+                        max_batch_size, max(1, len(split_values) // max_batch_size)
+                    )
+                    batch_size = min(len(split_values), max_batch_size)
 
                     # Setup progress bar for split parameter requests
-                    num_batches = ceil(len(split_values) / batch_size)
                     pbar_message = f"Retrieving {len(split_values)} {split_param} values in {num_batches} batches"
                     pbar = (
                         tqdm(
@@ -998,7 +1000,7 @@ def _submit_requests(  # noqa
                 desc=pbar_message,
                 total=num_docs_needed,
             )
-            if not self.mute_progress_bars
+            if not self.mute_progress_bars and total_num_docs > 0
             else None
         )