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 <[email protected]>" ]
 readme = "README.md"