Bugfix/ffs rate limit handler (#425)

* fix and changelog

* handle no retry-after header case

* handle no retry-after header case

* fix failing tests

* refactor to avoid complicating test mocking

* add fake session to mock connections to fix tests

* add fake session to mock connections to fix tests

* style

CHANGELOG.md

@@ -55,6 +55,7 @@ how a consumer would use the library (e.g. adding unit tests, updating documenta
 ### Fixed
 
 - Bug where attempting to restore from an empty archive would throw a confusing `TypeError`, we now raise appropriate `Py42ArchiveFileNotFoundError`.
+- Py42 now automatically sleeps and retries file event search queries if a 429 (too many requests) error is hit.
 
 ## 1.22.0 - 2022-04-01
 

src/py42/services/fileevent.py

@@ -1,18 +1,48 @@
 import json
 from warnings import warn
 
+from requests.adapters import HTTPAdapter
+from urllib3 import Retry
+
+import py42.settings.debug as debug
 from py42.exceptions import Py42BadRequestError
 from py42.exceptions import Py42InvalidPageTokenError
 from py42.services import BaseService
 
 
+class FFSQueryRetryStrategy(Retry):
+    """The forensic search service helpfully responds with a 'retry-after' header, telling us how long until the rate
+    limiter is reset. We subclass :class:`urllib3.Retry` just to add a bit of logging so the user can tell why the
+    request might look like it's hanging.
+    """
+
+    def get_retry_after(self, response):
+        retry_after = super().get_retry_after(response)
+        if retry_after is not None:
+            debug.logger.info(
+                f"Forensic search rate limit hit, retrying after: {int(retry_after)} seconds."
+            )
+        return retry_after
+
+    def get_backoff_time(self):
+        backoff_time = super().get_backoff_time()
+        debug.logger.info(
+            f"Forensic search rate limit hit, retrying after: {backoff_time} seconds."
+        )
+        return backoff_time
+
+
 class FileEventService(BaseService):
     """A service for searching file events.
 
     See the :ref:`Executing Searches User Guide <anchor_search_file_events>` to learn more about how
     to construct a query.
     """
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._retry_adapter_mounted = False
+
     def search(self, query):
         """Searches for file events matching the query criteria.
         `REST Documentation <https://developer.code42.com/api/#operation/searchEventsUsingPOST>`__
@@ -28,6 +58,8 @@ def search(self, query):
         Returns:
             :class:`py42.response.Py42Response`: A response containing the query results.
         """
+        self._mount_retry_adapter()
+
         # if string query
         if isinstance(query, str):
             query = json.loads(query)
@@ -65,5 +97,31 @@ def get_file_location_detail_by_sha256(self, checksum):
         Returns:
             :class:`py42.response.Py42Response`: A response containing file details.
         """
+        self._mount_retry_adapter()
+
         uri = "/forensic-search/queryservice/api/v1/filelocations"
         return self._connection.get(uri, params={"sha256": checksum})
+
+    def _mount_retry_adapter(self):
+        """Sets custom Retry strategy for FFS url requests to gracefully handle being rate-limited on FFS queries."""
+        if not self._retry_adapter_mounted:
+            retry_strategy = FFSQueryRetryStrategy(
+                status=3,  # retry up to 3 times
+                backoff_factor=5,  # if `retry-after` header isn't present, use 5 second exponential backoff
+                allowed_methods=[
+                    "POST"
+                ],  # POST isn't a default allowed method due to it usually modifying resources
+                status_forcelist=[
+                    429
+                ],  # this only handles 429 errors, it won't retry on 5xx
+            )
+            file_event_adapter = HTTPAdapter(
+                pool_connections=200,
+                pool_maxsize=4,
+                pool_block=True,
+                max_retries=retry_strategy,
+            )
+            self._connection._session.mount(
+                self._connection.host_address, file_event_adapter
+            )
+            self._retry_adapter_mounted = True

tests/clients/test_securitydata.py

@@ -1,4 +1,5 @@
 import pytest
+import requests
 from tests.conftest import create_mock_response
 
 from py42.clients.securitydata import SecurityDataClient
@@ -155,7 +156,9 @@
 class TestSecurityClient:
     @pytest.fixture
     def connection(self, mocker):
-        return mocker.MagicMock(spec=Connection)
+        connection = mocker.MagicMock(spec=Connection)
+        connection._session = mocker.MagicMock(spec=requests.Session)
+        return connection
 
     @pytest.fixture
     def storage_service_factory(self, mocker):

tests/conftest.py

@@ -182,6 +182,7 @@ def unicode_query_filter():
 @pytest.fixture
 def mock_connection(mocker):
     connection = mocker.MagicMock(spec=Connection)
+    connection._session = mocker.MagicMock(spec=Session)
     connection.headers = {}
 
     return connection

tests/services/test_fileevent.py

@@ -1,6 +1,7 @@
 import json
 
 import pytest
+import requests
 from tests.conftest import create_mock_error
 
 from py42.exceptions import Py42BadRequestError
@@ -39,7 +40,9 @@ def mock_invalid_page_token_connection(mocker, connection):
 class TestFileEventService:
     @pytest.fixture
     def connection(self, mocker):
-        return mocker.MagicMock(spec=Connection)
+        connection = mocker.MagicMock(spec=Connection)
+        connection._session = mocker.MagicMock(spec=requests.Session)
+        return connection
 
     def test_search_calls_post_with_uri_and_query(
         self, connection, successful_response