Merge pull request #115 from CrowdStrike/online-state-batch-fix

Online State Endpoints Batch Data Size Fix

caracara/common/__init__.py

@@ -1,13 +1,13 @@
 """Caracara: Common functions and data imports."""
 __all__ = [
-    "DATA_BATCH_SIZE",
+    "DEFAULT_DATA_BATCH_SIZE",
     "FalconApiModule",
     "Policy",
     "SCROLL_BATCH_SIZE",
     "user_agent_string",
 ]
 
-from caracara.common.constants import SCROLL_BATCH_SIZE, DATA_BATCH_SIZE
+from caracara.common.constants import SCROLL_BATCH_SIZE, DEFAULT_DATA_BATCH_SIZE
 from caracara.common.policy_wrapper import Policy
 from caracara.common.meta import user_agent_string
 from caracara.common.module import FalconApiModule

caracara/common/batching.py

@@ -16,7 +16,7 @@
 from threading import current_thread
 from typing import Callable, Dict, List, Tuple
 
-from caracara.common.constants import DATA_BATCH_SIZE
+from caracara.common.constants import DEFAULT_DATA_BATCH_SIZE
 
 
 BATCH_LOGGER = logging.getLogger(__name__)
@@ -37,6 +37,7 @@ def batch_data_pull_threads() -> int:
 def batch_get_data(
     lookup_ids: str,
     func: Callable[[object, List[str], str], Dict],
+    data_batch_size: int = DEFAULT_DATA_BATCH_SIZE,
 ) -> Dict[str, Dict]:
     """Retrieve details for the list of AIDs provided.
 
@@ -57,8 +58,8 @@ def batch_get_data(
     BATCH_LOGGER.info("Batch data retrieval for %s (%d items)", func.__name__, len(lookup_ids))
     BATCH_LOGGER.debug(str(lookup_ids))
 
-    # Divide the list of item IDs into a list of lists, each of size DATA_BATCH_SIZE
-    batches = [lookup_ids[i:i+DATA_BATCH_SIZE] for i in range(0, len(lookup_ids), DATA_BATCH_SIZE)]
+    # Divide the list of item IDs into a list of lists, each of size data_batch_size
+    batches = [lookup_ids[i:i+data_batch_size] for i in range(0, len(lookup_ids), data_batch_size)]
     BATCH_LOGGER.info("Divided the item IDs into %d batches", len(batches))
 
     threads = batch_data_pull_threads()

caracara/common/constants.py

@@ -2,7 +2,10 @@
 from enum import Enum, EnumMeta
 
 # Batch size of data downloaded via a multi-threaded data pull
-DATA_BATCH_SIZE = 500
+DEFAULT_DATA_BATCH_SIZE = 500
+
+# Batch size of data downloaded via a multi-threaded data pull from the online state endpoint
+ONLINE_STATE_DATA_BATCH_SIZE = 100
 
 # Default pagination limit
 PAGINATION_LIMIT = 100

caracara/modules/hosts/_online_state.py

@@ -12,7 +12,7 @@
 )
 
 from caracara.common.batching import batch_get_data
-from caracara.common.constants import OnlineState
+from caracara.common.constants import OnlineState, ONLINE_STATE_DATA_BATCH_SIZE
 from caracara.common.exceptions import InvalidOnlineState
 
 if TYPE_CHECKING:
@@ -56,7 +56,11 @@ def get_online_state(
     dict: A dictionary containing online state details for every device listed.
     """
     self.logger.info("Obtaining online state data for %s devices", len(device_ids))
-    device_online_state_data = batch_get_data(device_ids, self.hosts_api.get_online_state)
+    device_online_state_data = batch_get_data(
+        device_ids,
+        self.hosts_api.get_online_state,
+        data_batch_size=ONLINE_STATE_DATA_BATCH_SIZE,
+    )
     return device_online_state_data
 
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "caracara"
-version = "0.5.0"
+version = "0.5.1"
 description = "The CrowdStrike Falcon Developer Toolkit"
 authors = [ "CrowdStrike <falconpy@crowdstrike.com>" ]
 readme = "README.md"