From deab7a5e2ca10e29f40b4fb191ee09bdb6dae4c3 Mon Sep 17 00:00:00 2001
From: vx3r <vx3r@127-0-0-1.fr>
Date: Mon, 21 Oct 2024 19:28:48 +0200
Subject: [PATCH] Add max retry parameter to the execution (#802)

Co-authored-by: Meroujan.Antonyan <meroujan.antonyan.external@axa.com>
Co-authored-by: Ian Hellen <ianhelle@microsoft.com>
---
 msticpy/data/drivers/cybereason_driver.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/msticpy/data/drivers/cybereason_driver.py b/msticpy/data/drivers/cybereason_driver.py
index 3901f7f1f..2750b010d 100644
--- a/msticpy/data/drivers/cybereason_driver.py
+++ b/msticpy/data/drivers/cybereason_driver.py
@@ -401,6 +401,7 @@ def __execute_query(
         page: int = 0,
         page_size: int = 2000,
         pagination_token: str = None,
+        max_retry: int = 3,
     ) -> Dict[str, Any]:
         """
         Run query with pagination enabled.
@@ -436,7 +437,8 @@ def __execute_query(
             headers = {}
         params = {"page": page, "itemsPerPage": page_size}
         status = None
-        while status != "SUCCESS":
+        cur_try = 0
+        while status != "SUCCESS" and cur_try < max_retry:
             response = self.client.post(
                 self.search_endpoint,
                 json={**body, **pagination},
@@ -446,6 +448,7 @@ def __execute_query(
             response.raise_for_status()
             json_result = response.json()
             status = json_result["status"]
+            cur_try += 1
         return json_result
 
     async def __run_threaded_queries(