feat: Improve retry on some HTTP errors

CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## 1.15.1 - 2024-010-04
+
+### Changed
+
+- Improve retry on some HTTP errors
+
 ## [1.15.0] - 2024-09-28
 
 ### Changed

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "sekoia-automation-sdk"
 
-version = "1.15.0"
+version = "1.15.1"
 description = "SDK to create Sekoia.io playbook modules"
 license = "MIT"
 readme = "README.md"

sekoia_automation/action.py

@@ -302,8 +302,8 @@ def log_request_error(self, url: str, arguments: dict, response: Response):
         )
         self.error(message)
 
-    def log_timeout_error(self, url: str, arguments: dict):
-        message = f"HTTP Request timeout: {url}"
+    def log_retry_error(self, url: str, arguments: dict):
+        message = f"HTTP Request failed after all retries: {url}"
         self.log(
             message,
             level="error",
@@ -349,15 +349,22 @@ def run(self, arguments) -> dict | None:
                     response: Response = requests.request(
                         self.verb, url, json=body, headers=headers, timeout=self.timeout
                     )
+                    if not response.ok:
+                        if (
+                            self.verb.lower() == "delete"
+                            and response.status_code == 404
+                            and attempt.retry_state.attempt_number > 1
+                        ):
+                            return None
+                        if 400 <= response.status_code < 500:
+                            self.log_request_error(url, arguments, response)
+                            return None
+                        response.raise_for_status()
         except RetryError:
-            self.log_timeout_error(url, arguments)
-            return None
-
-        if not response.ok:
-            self.log_request_error(url, arguments, response)
+            self.log_retry_error(url, arguments)
             return None
 
         return response.json() if response.status_code != 204 else None
 
     def _wait_param(self) -> wait_base:
-        return wait_exponential(multiplier=2, min=1, max=10)
+        return wait_exponential(multiplier=2, min=2, max=300)

sekoia_automation/aio/connector.py

@@ -91,10 +91,7 @@ async def session(cls) -> AsyncGenerator[ClientSession, None]:  # pragma: no cov
         Returns:
             ClientSession:
         """
-        if cls._session is None:
-            cls._session = ClientSession()
-
-        async with cls.get_rate_limiter():
+        async with ClientSession() as cls._session, cls.get_rate_limiter():
             yield cls._session
 
     async def _async_send_chunk(

sekoia_automation/aio/helpers/http/http_client.py

@@ -81,11 +81,9 @@ async def session(self) -> AsyncGenerator[ClientSession, None]:
         Yields:
             AsyncGenerator[ClientSession, None]:
         """
-        if self._session is None:
-            self._session = ClientSession()
-
-        if self._rate_limiter:
-            async with self._rate_limiter:
+        async with ClientSession() as self._session:
+            if self._rate_limiter:
+                async with self._rate_limiter:
+                    yield self._session
+            else:
                 yield self._session
-        else:
-            yield self._session

sekoia_automation/http/aio/http_client.py

@@ -45,14 +45,12 @@ async def session(self) -> AsyncGenerator[ClientSession, None]:
         Yields:
             AsyncGenerator[ClientSession, None]:
         """
-        if self._session is None:
-            self._session = ClientSession()
-
-        if self._rate_limiter:
-            async with self._rate_limiter:
+        async with ClientSession() as self._session:
+            if self._rate_limiter:
+                async with self._rate_limiter:
+                    yield self._session
+            else:
                 yield self._session
-        else:
-            yield self._session
 
     @asynccontextmanager
     async def get(

sekoia_automation/module.py

@@ -11,7 +11,7 @@
 import sentry_sdk
 from botocore.exceptions import ClientError
 from pydantic import BaseModel
-from requests import HTTPError, Response
+from requests import RequestException, Response
 
 from sekoia_automation.config import load_config
 from sekoia_automation.exceptions import (
@@ -441,19 +441,22 @@ def _send_request(self, data: dict, verb: str = "POST", attempt=1) -> Response:
             )
             response.raise_for_status()
             return response
-        except HTTPError as exception:
-            self._log_request_error(exception)
+        except (RequestException, OSError) as exception:
+            if isinstance(exception, RequestException):
+                self._log_request_error(exception)
             if attempt == 10:
                 status_code = (
                     exception.response.status_code
-                    if isinstance(exception.response, Response)
+                    if isinstance(exception, RequestException)
+                    and isinstance(exception.response, Response)
                     else 500
                 )
                 raise SendEventError(
                     "Impossible to send event to Sekoia.io API", status_code=status_code
                 )
             if (
-                isinstance(exception.response, Response)
+                isinstance(exception, RequestException)
+                and isinstance(exception.response, Response)
                 and 400 <= exception.response.status_code < 500
             ):
                 raise SendEventError(
@@ -463,7 +466,7 @@ def _send_request(self, data: dict, verb: str = "POST", attempt=1) -> Response:
             time.sleep(self._wait_exponent_base**attempt)
             return self._send_request(data, verb, attempt + 1)
 
-    def _log_request_error(self, exception: HTTPError):
+    def _log_request_error(self, exception: RequestException):
         context: dict[str, Any] = {}
         if exception.response:
             response: Response = exception.response

tests/aio/test_connector.py

@@ -76,17 +76,14 @@ async def test_async_connector_client_session(async_connector: DummyAsyncConnect
     assert async_connector._session is None
     assert other_instance._session is None
 
-    async with async_connector.session() as session_1:
-        async with other_instance.session() as session_2:
-            assert session_1 == session_2
-
-            assert async_connector._rate_limiter is not None and isinstance(
-                async_connector._rate_limiter, AsyncLimiter
-            )
+    async with async_connector.session():
+        assert async_connector._rate_limiter is not None and isinstance(
+            async_connector._rate_limiter, AsyncLimiter
+        )
 
-            assert other_instance._rate_limiter is not None and isinstance(
-                other_instance._rate_limiter, AsyncLimiter
-            )
+        assert other_instance._rate_limiter is not None and isinstance(
+            other_instance._rate_limiter, AsyncLimiter
+        )
 
     DummyAsyncConnector.set_rate_limiter(None)
     other_instance.set_client_session(None)

tests/test_action.py

@@ -39,12 +39,10 @@ def test_action_logs(capsys):
     assert action.logs[1]["level"] == "info"
     assert action.logs[2]["level"] == "warning"
     assert action.logs[3]["level"] == "error"
-    assert action.logs[4]["level"] == "warning"
     assert action.logs[0]["message"] == "message1"
     assert action.logs[1]["message"] == "message2"
     assert action.logs[2]["message"] == "message3"
     assert action.logs[3]["message"] == "message4"
-    assert action.logs[4]["message"] == "message5"
 
 
 def test_action_outputs():
@@ -239,9 +237,9 @@ def test_action_json_result_same_as_argument():
 
 
 def test_generic_api_action(storage):
-    def init_action():
+    def init_action(verb: str = "get"):
         action = GenericAPIAction(data_path=storage)
-        action.verb = "get"
+        action.verb = verb
         action.endpoint = "resource/{uuid}/count"
         action.query_parameters = ["param"]
         action.module.configuration = {"base_url": "http://base_url/"}
@@ -321,9 +319,28 @@ def init_action():
         mock.get("http://base_url/resource/fake_uuid/count", status_code=500, json={})
         results: dict = action.run({"uuid": "fake_uuid", "param": "number"})
 
+        assert results is None
+        assert mock.call_count == 10
+
+    action = init_action()
+    with requests_mock.Mocker() as mock:
+        mock.get("http://base_url/resource/fake_uuid/count", status_code=400, json={})
+        results: dict = action.run({"uuid": "fake_uuid", "param": "number"})
+
         assert results is None
         assert mock.call_count == 1
 
+    action = init_action(verb="delete")
+    with requests_mock.Mocker() as mock:
+        mock.delete(
+            "http://base_url/resource/fake_uuid/count",
+            [{"status_code": 503}, {"status_code": 404}],
+        )
+        results: dict = action.run({"uuid": "fake_uuid", "param": "number"})
+
+        assert results is None
+        assert mock.call_count == 2
+
     # timeout
     action = init_action()
     arguments = {"uuid": "fake_uuid", "param": "number"}