Configurable batch size for JIRA._fetch_pages() and dependant methods

constraints.txt

@@ -83,6 +83,8 @@ packaging==21.3
     #   pytest
     #   pytest-sugar
     #   sphinx
+parameterized==0.8.1
+    # via jira (setup.cfg)
 parso==0.8.3
     # via jedi
 pickleshare==0.7.5

jira/client.py

@@ -366,6 +366,10 @@ class JIRA:
             # 'Expires': 'Thu, 01 Jan 1970 00:00:00 GMT'
             "X-Atlassian-Token": "no-check",
         },
+        "default_batch_size": {
+            Resource: 100,
+            Issue: 500,
+        },
     }
 
     checked_version = False
@@ -393,6 +397,7 @@ def __init__(
         proxies: Any = None,
         timeout: Optional[Union[Union[float, int], Tuple[float, float]]] = None,
         auth: Tuple[str, str] = None,
+        default_batch_sizes: Optional[Dict[ResourceType, Optional[int]]] = None,
     ):
         """Construct a Jira client instance.
 
@@ -472,7 +477,6 @@ def __init__(
         # force a copy of the tuple to be used in __del__() because
         # sys.version_info could have already been deleted in __del__()
         self.sys_version_info = tuple(sys.version_info)
-
         if options is None:
             options = {}
             if server and isinstance(server, dict):
@@ -494,6 +498,9 @@ def __init__(
 
         self._options: Dict[str, Any] = copy.copy(JIRA.DEFAULT_OPTIONS)
 
+        if default_batch_sizes:
+            self._options["default_batch_size"].update(default_batch_sizes)
+
         if "headers" in options:
             headers = copy.copy(options["headers"])
             del options["headers"]
@@ -681,7 +688,6 @@ def _fetch_pages(
         request_path: str,
         startAt: int = 0,
         maxResults: int = 50,
-        batch_size: Optional[int] = None,
         params: Dict[str, Any] = None,
         base: str = JIRA_BASE_URL,
     ) -> ResultList[ResourceType]:
@@ -695,9 +701,6 @@ def _fetch_pages(
             startAt (int): index of the first record to be fetched. (Default: 0)
             maxResults (int): Maximum number of items to return.
               If maxResults evaluates as False, it will try to get all items in batches. (Default:50)
-            batch_size (Optional[int]): Size of batches, only used if maxResults evaluates as False.
-              If not specified, the JIRA-instance's default size will be used. Strongly recommended to use a value over
-              100 for larger requests.
             params (Dict[str, Any]): Params to be used in all requests. Should not contain startAt and maxResults,
               as they will be added for each request created from this function.
             base (str): base URL to use for the requests.
@@ -720,7 +723,7 @@ def _fetch_pages(
             page_params["startAt"] = startAt
         if maxResults:
             page_params["maxResults"] = maxResults
-        elif batch_size is not None:
+        elif batch_size := self._get_batch_size(item_type):
             page_params["maxResults"] = batch_size
 
         resource = self._get_json(request_path, params=page_params, base=base)
@@ -827,6 +830,16 @@ def _get_items_from_page(
             # improving the error text so we know why it happened
             raise KeyError(str(e) + " : " + json.dumps(resource))
 
+    def _get_batch_size(self, item_type: Type[ResourceType]) -> Optional[int]:
+        batch_sizes: Dict[Type[Resource], Optional[int]] = self._options[
+            "default_batch_size"
+        ]
+        try:
+            item_type_batch_size = batch_sizes[item_type]
+        except KeyError:
+            item_type_batch_size = batch_sizes.get(Resource, None)
+        return item_type_batch_size
+
     # Information about this client
 
     def client_info(self) -> str:
@@ -1139,17 +1152,14 @@ def custom_field_option(self, id: str) -> CustomFieldOption:
     # Dashboards
 
     def dashboards(
-        self, filter=None, startAt=0, maxResults=20, batch_size: Optional[int] = None
+        self, filter=None, startAt=0, maxResults=20
     ) -> ResultList[Dashboard]:
         """Return a ResultList of Dashboard resources and a ``total`` count.
 
         Args:
             filter (Optional[str]): either "favourite" or "my", the type of dashboards to return
             startAt (int): index of the first dashboard to return (Default: 0)
             maxResults (int): maximum number of dashboards to return. If maxResults evaluates as False, it will try to get all items in batches. (Default: 20)
-            batch_size (Optional[int]): Size of batches, only used if maxResults evaluates as False.
-              If not specified, the JIRA-instance's default size will be used. Strongly recommended to use a value over
-              100 for larger requests.
 
         Returns:
             ResultList
@@ -1163,7 +1173,6 @@ def dashboards(
             "dashboard",
             startAt,
             maxResults,
-            batch_size,
             params,
         )
 
@@ -2894,7 +2903,6 @@ def search_issues(
         fields: Optional[Union[str, List[str]]] = "*all",
         expand: Optional[str] = None,
         json_result: bool = False,
-        batch_size: Optional[int] = None,
     ) -> Union[List[Dict[str, Any]], ResultList[Issue]]:
         """Get a :class:`~jira.client.ResultList` of issue Resources matching a JQL search string.
 
@@ -2910,9 +2918,6 @@ def search_issues(
             expand (Optional[str]): extra information to fetch inside each resource
             json_result (bool): JSON response will be returned when this parameter is set to True.
               Otherwise, :class:`~jira.client.ResultList` will be returned.
-            batch_size (Optional[int]): Size of batches, only used if maxResults evaluates as False.
-              If not specified, the JIRA-instance's default size will be used. Strongly recommended to use a value over
-              100 for larger requests.
 
         Returns:
             Union[Dict,ResultList]: Dict if ``json_result=True``
@@ -2950,7 +2955,7 @@ def search_issues(
             return r_json
 
         issues = self._fetch_pages(
-            Issue, "issues", "search", startAt, maxResults, batch_size, search_params
+            Issue, "issues", "search", startAt, maxResults, search_params
         )
 
         if untranslate:
@@ -3077,7 +3082,6 @@ def search_assignable_users_for_projects(
         projectKeys: str,
         startAt: int = 0,
         maxResults: int = 50,
-        batch_size: Optional[int] = None,
     ) -> ResultList:
         """Get a list of user Resources that match the search string and can be assigned issues for projects.
 
@@ -3087,9 +3091,6 @@ def search_assignable_users_for_projects(
             startAt (int): Index of the first user to return (Default: 0)
             maxResults (int): Maximum number of users to return.
               If maxResults evaluates as False, it will try to get all users in batches. (Default: 50)
-            batch_size (Optional[int]): Size of batches, only used if maxResults evaluates as False.
-              If not specified, the JIRA-instance's default size will be used. Strongly recommended to use a value over
-              100 for larger requests.
 
         Returns:
             ResultList
@@ -3102,7 +3103,6 @@ def search_assignable_users_for_projects(
             "user/assignable/multiProjectSearch",
             startAt,
             maxResults,
-            batch_size,
             params,
         )
 
@@ -3115,7 +3115,6 @@ def search_assignable_users_for_issues(
         startAt: int = 0,
         maxResults: int = 50,
         query: Optional[str] = None,
-        batch_size: Optional[int] = None,
     ):
         """Get a list of user Resources that match the search string for assigning or creating issues.
         "username" query parameter is deprecated in Jira Cloud; the expected parameter now is "query", which can just be
@@ -3136,9 +3135,6 @@ def search_assignable_users_for_issues(
             maxResults (int): maximum number of users to return.
               If maxResults evaluates as False, it will try to get all items in batches. (Default: 50)
             query (Optional[str]): Search term. It can just be the email.
-            batch_size (Optional[int]): Size of batches, only used if maxResults evaluates as False.
-              If not specified, the JIRA-instance's default size will be used. Strongly recommended to use a value over
-              100 for larger requests.
 
         Returns:
             ResultList
@@ -3165,7 +3161,6 @@ def search_assignable_users_for_issues(
             "user/assignable/search",
             startAt,
             maxResults,
-            batch_size,
             params,
         )
 
@@ -3288,7 +3283,6 @@ def search_users(
         includeActive: bool = True,
         includeInactive: bool = False,
         query: Optional[str] = None,
-        batch_size: Optional[int] = None,
     ) -> ResultList[User]:
         """Get a list of user Resources that match the specified search string.
         "username" query parameter is deprecated in Jira Cloud; the expected parameter now is "query", which can just be the full
@@ -3302,9 +3296,6 @@ def search_users(
             includeActive (bool): If true, then active users are included in the results. (Default: True)
             includeInactive (bool): If true, then inactive users are included in the results. (Default: False)
             query (Optional[str]): Search term. It can just be the email.
-            batch_size (Optional[int]): Size of batches, only used if maxResults evaluates as False.
-              If not specified, the JIRA-instance's default size will be used. Strongly recommended to use a value over
-              100 for larger requests.
 
         Returns:
             ResultList[User]
@@ -3319,9 +3310,7 @@ def search_users(
             "includeInactive": includeInactive,
         }
 
-        return self._fetch_pages(
-            User, None, "user/search", startAt, maxResults, batch_size, params
-        )
+        return self._fetch_pages(User, None, "user/search", startAt, maxResults, params)
 
     def search_allowed_users_for_issue(
         self,
@@ -3330,7 +3319,6 @@ def search_allowed_users_for_issue(
         projectKey: str = None,
         startAt: int = 0,
         maxResults: int = 50,
-        batch_size: Optional[int] = None,
     ) -> ResultList:
         """Get a list of user Resources that match a username string and have browse permission for the issue or project.
 
@@ -3341,9 +3329,6 @@ def search_allowed_users_for_issue(
             startAt (int): index of the first user to return. (Default: 0)
             maxResults (int): maximum number of users to return.
               If maxResults evaluates as False, it will try to get all items in batches. (Default: 50)
-            batch_size (Optional[int]): Size of batches, only used if maxResults evaluates as False.
-              If not specified, the JIRA-instance's default size will be used. Strongly recommended to use a value over
-              100 for larger requests.
 
         Returns:
             ResultList
@@ -3354,7 +3339,7 @@ def search_allowed_users_for_issue(
         if projectKey is not None:
             params["projectKey"] = projectKey
         return self._fetch_pages(
-            User, None, "user/viewissue/search", startAt, maxResults, batch_size, params
+            User, None, "user/viewissue/search", startAt, maxResults, params
         )
 
     # Versions
@@ -4516,7 +4501,6 @@ def boards(
         type: str = None,
         name: str = None,
         projectKeyOrID=None,
-        batch_size: Optional[int] = None,
     ) -> ResultList[Board]:
         """Get a list of board resources.
 
@@ -4526,9 +4510,6 @@ def boards(
             type: Filters results to boards of the specified type. Valid values: scrum, kanban.
             name: Filters results to boards that match or partially match the specified name.
             projectKeyOrID: Filters results to boards that match the specified project key or ID.
-            batch_size (Optional[int]): Size of batches, only used if maxResults evaluates as False.
-              If not specified, the JIRA-instance's default size will be used. Strongly recommended to use a value over
-              100 for larger requests.
 
         Returns:
             ResultList[Board]
@@ -4547,7 +4528,6 @@ def boards(
             "board",
             startAt,
             maxResults,
-            batch_size,
             params,
             base=self.AGILE_BASE_URL,
         )
@@ -4560,7 +4540,6 @@ def sprints(
         startAt: int = 0,
         maxResults: int = 50,
         state: str = None,
-        batch_size: Optional[int] = None,
     ) -> ResultList[Sprint]:
         """Get a list of sprint Resources.
 
@@ -4571,9 +4550,6 @@ def sprints(
             maxResults (int): the maximum number of sprints to return
             state (str): Filters results to sprints in specified states. Valid values: `future`, `active`, `closed`.
               You can define multiple states separated by commas
-            batch_size (Optional[int]): Size of batches, only used if maxResults evaluates as False.
-              If not specified, the JIRA-instance's default size will be used. Strongly recommended to use a value over
-              100 for larger requests.
 
         Returns:
             ResultList[Sprint]: List of sprints.
@@ -4591,7 +4567,6 @@ def sprints(
             f"board/{board_id}/sprint",
             startAt,
             maxResults,
-            batch_size,
             params,
             self.AGILE_BASE_URL,
         )

tests/tests.py

@@ -12,7 +12,7 @@
 import os
 import pickle
 from time import sleep
-from typing import Optional, cast
+from typing import cast
 from unittest import mock
 
 import pytest
@@ -21,7 +21,7 @@
 
 from jira import JIRA, Issue, JIRAError
 from jira.client import ResultList
-from jira.resources import cls_for_resource
+from jira.resources import Resource, cls_for_resource
 from tests.conftest import JiraTestCase, rndpassword
 
 LOGGER = logging.getLogger(__name__)
@@ -302,20 +302,23 @@ def setUp(self):
             (
                 0,
                 26,
-                None,
+                {Issue: None},
                 False,
-            ),  # original behaviour, fetch all with jira's original return size (here 10)
-            (0, 26, 20, False),  # set batch size to 20
-            (5, 26, 20, False),  # test start_at
-            (5, 26, 20, 50),  # test maxResults set (one request)
+            ),  # original behaviour, fetch all with jira's original return size
+            (0, 26, {Issue: 20}, False),  # set batch size to 20
+            (5, 26, {Issue: 20}, False),  # test start_at
+            (5, 26, {Issue: 20}, 50),  # test maxResults set (one request)
         ]
     )
     def test_fetch_pages(
-        self, start_at: int, total: int, batch_size: Optional[int], max_results: int
+        self, start_at: int, total: int, default_batch_sizes: dict, max_results: int
     ):
         """Tests that the JIRA._fetch_pages method works as expected."""
         params = {"startAt": 0}
-        batch_size = batch_size or 10
+        self.jira._options["default_batch_size"] = default_batch_sizes
+        batch_size = (
+            self.jira._get_batch_size(Issue) or 10
+        )  # 10 -> mimicked JIRA-backend default if we did not specify it
         expected_calls = _calculate_calls_for_fetch_pages(
             "https://jira.atlassian.com/rest/api/2/search",
             start_at,
@@ -356,7 +359,7 @@ def test_fetch_pages(
         self.jira._session.close()
         self.jira._session = mock_session
         items = self.jira._fetch_pages(
-            Issue, "issues", "search", start_at, max_results, batch_size, params=params
+            Issue, "issues", "search", start_at, max_results, params=params
         )
 
         actual_calls = [[kall[1], kall[2]] for kall in self.jira._session.method_calls]
@@ -368,6 +371,25 @@ def test_fetch_pages(
         )
 
 
+@parameterized.expand(
+    [
+        ({Resource: 1, Issue: 2}, Issue, 2),
+        ({Resource: 1, Issue: 2}, Resource, 1),
+        ({Resource: 1, Issue: None}, Issue, None),
+        ({Resource: 1}, Issue, 1),
+        ({Resource: 1}, Issue, 1),
+    ]
+)
+def test_get_batch_size(default_batch_sizes, item_type, expected):
+    class BatchSizeMock:
+        def __init__(self, batch_sizes):
+            self._options = {"default_batch_size": batch_sizes}
+
+    batch_size_mock = BatchSizeMock(default_batch_sizes)
+
+    assert JIRA._get_batch_size(batch_size_mock, item_type) == expected
+
+
 def _create_issue_result_json(issue_id, summary, key, **kwargs):
     """Returns a minimal json object for an issue."""
     return {