Improve handling for errors while fetching full events (#94)

Co-authored-by: Mark Kasaboski <mark.kasaboski@gmail.com>
Co-authored-by: Xavier Lavallee <xavier.lavallee@gmail.com>

Author: 3 people

packages/flare/bin/flare.py

@@ -46,12 +46,18 @@ def get_flare_api_client(
 
 
 class FlareAPI(AuthBase):
-    def __init__(self, *, api_key: str, tenant_id: Optional[int] = None) -> None:
+    def __init__(
+        self,
+        *,
+        api_key: str,
+        tenant_id: Optional[int] = None,
+        logger: Logger,
+    ) -> None:
         self.flare_client = get_flare_api_client(
             api_key=api_key,
             tenant_id=tenant_id,
         )
-        self.logger = Logger(class_name=__file__)
+        self.logger = logger
 
     def fetch_feed_events(
         self,
@@ -76,7 +82,7 @@ def fetch_feed_events(
                     event = self._fetch_full_event_from_uid(
                         uid=event["metadata"]["uid"]
                     )
-                    time.sleep(1)
+                    time.sleep(1)  # Don't hit rate limit
                 yield (event, next_token)
 
     def _fetch_event_feed_metadata(
@@ -115,10 +121,23 @@ def _fetch_event_feed_metadata(
             time.sleep(1)
 
     def _fetch_full_event_from_uid(self, *, uid: str) -> dict:
-        event_response = self.flare_client.get(url=f"/firework/v2/activities/{uid}")
-        event = event_response.json()["activity"]
-        self.logger.debug(event)
-        return event
+        number_of_retries = 3
+        for current_try in range(number_of_retries):
+            try:
+                event_response = self.flare_client.get(
+                    url=f"/firework/v2/activities/{uid}"
+                )
+                event_response.raise_for_status()
+            except Exception as e:
+                time.sleep(1)
+                self.logger.info(
+                    f"Failed to fetch event {current_try + 1}/{number_of_retries} retries: {e}"
+                )
+                continue
+            return event_response.json()["activity"]
+        raise Exception(
+            f"failed to fetch full event data for {uid} after {number_of_retries} tries"
+        )
 
     def fetch_api_key_validation(self) -> requests.Response:
         return self.flare_client.get(

packages/flare/bin/flare_external_requests.py

@@ -21,13 +21,14 @@
 
 class FlareValidateApiKey(splunk.rest.BaseRestHandler):
     def handle_POST(self) -> None:
+        logger = Logger(class_name=__file__)
         payload = self.request["payload"]
         params = parse.parse_qs(payload)
 
         if "apiKey" not in params:
             raise Exception("API Key is required")
 
-        flare_api = FlareAPI(api_key=params["apiKey"][0])
+        flare_api = FlareAPI(api_key=params["apiKey"][0], logger=logger)
         flare_api.fetch_api_key_validation()
         self.response.setHeader("Content-Type", "application/json")
         self.response.write(json.dumps({}))
@@ -42,7 +43,7 @@ def handle_POST(self) -> None:
         if "apiKey" not in params:
             raise Exception("API Key is required")
 
-        flare_api = FlareAPI(api_key=params["apiKey"][0])
+        flare_api = FlareAPI(api_key=params["apiKey"][0], logger=logger)
         response = flare_api.fetch_tenants()
         response_json = response.json()
         logger.debug(f"FlareUserTenants: {response_json}")
@@ -59,7 +60,7 @@ def handle_POST(self) -> None:
         if "apiKey" not in params:
             raise Exception("API Key is required")
 
-        flare_api = FlareAPI(api_key=params["apiKey"][0])
+        flare_api = FlareAPI(api_key=params["apiKey"][0], logger=logger)
         response = flare_api.fetch_filters_severity()
         response_json = response.json()
         logger.debug(f"FlareSeverityFilters: {response_json}")
@@ -76,7 +77,7 @@ def handle_POST(self) -> None:
         if "apiKey" not in params:
             raise Exception("API Key is required")
 
-        flare_api = FlareAPI(api_key=params["apiKey"][0])
+        flare_api = FlareAPI(api_key=params["apiKey"][0], logger=logger)
         response = flare_api.fetch_filters_source_type()
         response_json = response.json()
         logger.debug(f"FlareSourceTypeFilters: {response_json}")

packages/flare/tests/bin/conftest.py

@@ -11,7 +11,11 @@
 
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../bin"))
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../bin/vendor"))
 from data_store import ConfigDataStore
+from flare import FlareAPI
+from logger import Logger
+from vendor.splunklib.client import StoragePasswords
 
 
 class FakeStoragePassword:
@@ -34,26 +38,29 @@ def clear_password(self) -> str:
         return self._state["clear_password"]
 
 
-class FakeStoragePasswords:
+class FakeStoragePasswords(StoragePasswords):
     def __init__(self, passwords: List[FakeStoragePassword]) -> None:
         self._passwords = passwords
 
     def list(self) -> List[FakeStoragePassword]:
         return self._passwords
 
 
-class FakeLogger:
+class FakeLogger(Logger):
     def __init__(self) -> None:
+        super().__init__(class_name="Logger")
         self.messages: List[str] = []
 
+        self._mock = mock.MagicMock(spec=Logger)
+
     def info(self, message: str) -> None:
         self.messages.append(f"INFO: {message}")
 
     def error(self, message: str) -> None:
         self.messages.append(f"ERROR: {message}")
 
 
-class FakeFlareAPI:
+class FakeFlareAPI(FlareAPI):
     def __init__(self, api_key: str, tenant_id: int) -> None:
         pass
 
@@ -92,7 +99,7 @@ def storage_passwords(request: pytest.FixtureRequest) -> FakeStoragePasswords:
 
 
 @pytest.fixture
-def logger() -> FakeLogger:
+def logger() -> Logger:
     return FakeLogger()
 
 
@@ -121,3 +128,9 @@ def data_store(mock_env: None) -> Generator[ConfigDataStore, None, None]:
         store = ConfigDataStore()
         store._commit = lambda: None
         yield store
+
+
+@pytest.fixture
+def disable_sleep() -> Generator[None, None, None]:
+    with mock.patch("time.sleep", return_value=None):
+        yield