Merge pull request #74 from SEKOIA-IO/feat/retry_storage

feat: Retry errors during storage access

CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [1.4.1] - 2023-09-13
+
+### Added
+
+- Add retry when accessing the storage
+
 ## [1.4.0] - 2023-09-09
 
 ### Added

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "sekoia-automation-sdk"
 
-version = "1.4.0"
+version = "1.4.1"
 description = "SDK to create Sekoia.io playbook modules"
 license = "MIT"
 readme = "README.md"
@@ -33,8 +33,8 @@ python = ">=3.10,<3.12"
 requests = "^2.25"
 sentry-sdk = "*"
 tenacity = "*"
-boto3 = "^1.26"
-s3path = "^0.4"
+boto3 = "^1.28"
+s3path = "^0.5"
 orjson = "^3.8"
 pydantic = "^1.10"
 typer = { extras = ["all"], version = "^0.7"}

sekoia_automation/module.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import sys
+import time
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Any, cast
@@ -17,7 +18,10 @@
     SendEventError,
 )
 from sekoia_automation.storage import get_data_path
-from sekoia_automation.utils import get_annotation_for, get_as_model
+from sekoia_automation.utils import (
+    get_annotation_for,
+    get_as_model,
+)
 
 
 class Module:
@@ -304,6 +308,8 @@ class ModuleItem(ABC):
     description: str | None = None
     results_model: type[BaseModel] | None = None
 
+    _http_error_base_sleep = 0.5
+
     def __init__(self, module: Module | None = None, data_path: Path | None = None):
         self.module: Module = module or Module()
 
@@ -391,7 +397,7 @@ def _send_request(self, data: dict, verb: str = "POST", attempt=1) -> Response:
             return response
         except HTTPError as exception:
             self._log_request_error(exception)
-            if attempt == 3:
+            if attempt == 5:
                 status_code = (
                     exception.response.status_code
                     if isinstance(exception.response, Response)
@@ -408,6 +414,7 @@ def _send_request(self, data: dict, verb: str = "POST", attempt=1) -> Response:
                     "Impossible to send event to Sekoia.io API",
                     status_code=exception.response.status_code,
                 )
+            time.sleep(attempt * self._http_error_base_sleep)
             return self._send_request(data, verb, attempt + 1)
 
     def _log_request_error(self, exception: HTTPError):

sekoia_automation/storage.py

@@ -7,14 +7,22 @@
 import orjson
 from botocore.config import Config
 from s3path import S3Path, register_configuration_parameter
+from tenacity import retry, stop_after_attempt, wait_exponential
 
 from sekoia_automation import constants
 from sekoia_automation.config import VOLUME_PATH, load_config
+from sekoia_automation.utils import capture_retry_error
 
 FilePath = Path | str
 
 
 @lru_cache
+@retry(
+    reraise=True,
+    wait=wait_exponential(max=6),
+    stop=stop_after_attempt(10),
+    retry_error_callback=capture_retry_error,
+)
 def get_s3_data_path() -> Path:
     ca_path, cert_path, key_path = _get_tls_client_credentials()
     config = Config()
@@ -133,6 +141,12 @@ def __init__(self, filepath: FilePath, data_path: Path | None = None):
 
         self._data: dict = {}
 
+    @retry(
+        reraise=True,
+        wait=wait_exponential(max=6),
+        stop=stop_after_attempt(10),
+        retry_error_callback=capture_retry_error,
+    )
     def load(self):
         if not self._data and self._filepath.is_file():
             with self._filepath.open("r") as fd:
@@ -142,16 +156,31 @@ def load(self):
                     # The content is not valid json
                     self._data = {}
 
+    @retry(
+        reraise=True,
+        wait=wait_exponential(max=6),
+        stop=stop_after_attempt(10),
+        retry_error_callback=capture_retry_error,
+    )
+    def dump(self):
+        with self._filepath.open("w") as out:
+            out.write(orjson.dumps(self._data).decode("utf-8"))
+
     def __enter__(self):
         self.load()
 
         return self._data
 
     def __exit__(self, _, __, ___):
-        with self._filepath.open("w") as out:
-            out.write(orjson.dumps(self._data).decode("utf-8"))
+        self.dump()
 
 
+@retry(
+    reraise=True,
+    wait=wait_exponential(max=6),
+    stop=stop_after_attempt(10),
+    retry_error_callback=capture_retry_error,
+)
 def temp_directory(data_path: Path | None = None) -> str:
     """Create a temporary directory inside data storage.
 
@@ -166,6 +195,12 @@ def temp_directory(data_path: Path | None = None) -> str:
     return name
 
 
+@retry(
+    reraise=True,
+    wait=wait_exponential(max=6),
+    stop=stop_after_attempt(10),
+    retry_error_callback=capture_retry_error,
+)
 def write(
     filename: str,
     content: dict | str,

sekoia_automation/trigger.py

@@ -70,6 +70,12 @@ def __init__(self, module: Module | None = None, data_path: Path | None = None):
         self._liveness_server = None
         self._exporter = None
 
+    @retry(
+        reraise=True,
+        wait=wait_exponential(max=10),
+        stop=stop_after_attempt(10),
+        retry_error_callback=capture_retry_error,
+    )
     def _get_secrets_from_server(self) -> dict[str, Any]:
         """Calls the API to fetch this trigger's secrets
 
@@ -278,6 +284,7 @@ def log(self, message: str, level: str = "info", *args, **kwargs) -> None:
             self._critical_log_sent = True
 
     @retry(
+        reraise=True,
         wait=wait_exponential(max=10),
         stop=stop_after_attempt(10),
         retry_error_callback=capture_retry_error,

tests/test_trigger.py

@@ -103,6 +103,7 @@ def test_send_event_4xx_error(mocked_trigger_logs):
 
 def test_send_event_too_many_failures(mocked_trigger_logs):
     trigger = DummyTrigger()
+    trigger._http_error_base_sleep = 0
 
     mocked_trigger_logs.post(
         "http://sekoia-playbooks/callback",