Merge pull request #1162 from SEKOIA-IO/lv/tehtris_save_checkpoint

Tehtris - save the progress

Tehtris/CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+## 2024-11-07 - 1.16.0
+
+### Changed
+
+- Save the progress
+
 ## 2024-10-30 - 1.15.3
 
 ### Fixed

Tehtris/manifest.json

@@ -29,7 +29,7 @@
   "name": "TEHTRIS",
   "uuid": "1528d749-d353-4e38-ab1b-6e01d7595569",
   "slug": "tehtris",
-  "version": "1.15.3",
+  "version": "1.16.0",
   "categories": [
     "Endpoint"
   ]

Tehtris/pyproject.toml

@@ -6,10 +6,10 @@ authors = []
 
 [tool.poetry.dependencies]
 python = ">=3.10,<3.12"
-sekoia-automation-sdk = "^1.13.0"
-orjson = "^3.7.7"
-python-dateutil = "^2.8.2"
-cachetools = "^5.4.0"
+sekoia-automation-sdk = "^1.17.2"
+orjson = "^3.10.11"
+python-dateutil = "^2.9.0.post0"
+cachetools = "^5.5.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "*"

Tehtris/tehtris_modules/trigger_tehtris_events.py

@@ -7,7 +7,9 @@
 import orjson
 from cachetools import Cache, LRUCache
 from dateutil.parser import isoparse
+from sekoia_automation.checkpoint import CheckpointDatetime
 from sekoia_automation.connector import Connector, DefaultConnectorConfiguration
+from sekoia_automation.storage import PersistentJSON
 
 from tehtris_modules import TehtrisModule
 from tehtris_modules.client import ApiClient
@@ -31,7 +33,12 @@ class TehtrisEventConnector(Connector):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.from_date = datetime.now(timezone.utc) - timedelta(minutes=1)
+
+        self.cursor = CheckpointDatetime(
+            path=self._data_path, start_at=timedelta(minutes=1), ignore_older_than=timedelta(hours=1)
+        )
+        self.from_date = self.cursor.offset
+
         self.fetch_events_limit = 100
         self.events_cache: Cache = LRUCache(maxsize=1000)  # TODO: is it enough to have 1000 event ids in cache?
 
@@ -40,11 +47,6 @@ def client(self):
         return ApiClient(self.module.configuration.apikey)
 
     def __fetch_next_events(self, from_date: datetime, offset: int):
-        # We don't retrieve messages older than one hour
-        one_hour_ago = datetime.now(timezone.utc) - timedelta(hours=1)
-        if from_date < one_hour_ago:
-            from_date = one_hour_ago
-
         params = {
             "fromDate": int(from_date.timestamp()),
             "limit": self.fetch_events_limit,
@@ -137,6 +139,10 @@ def fetch_events(self) -> Generator[list, None, None]:
         current_lag: int = 0
         if most_recent_date_seen > self.from_date:
             self.from_date = most_recent_date_seen
+
+            # save in context the most recent date seen
+            self.cursor.offset = self.from_date
+
             delta_time = datetime.now(timezone.utc) - most_recent_date_seen
             current_lag = int(delta_time.total_seconds())
 
@@ -192,5 +198,5 @@ def next_batch(self):
             time.sleep(delta_sleep)
 
     def run(self):
-        while True:
+        while self.running:
             self.next_batch()

Tehtris/tests/test_tehtris_event_trigger.py

@@ -16,7 +16,7 @@ def fake_time():
 
 @pytest.fixture
 def patch_datetime_now(fake_time):
-    with patch("tehtris_modules.trigger_tehtris_events.datetime") as mock_datetime:
+    with patch("sekoia_automation.checkpoint.datetime") as mock_datetime:
         mock_datetime.now.return_value = fake_time
         mock_datetime.side_effect = lambda *args, **kw: datetime(*args, **kw)
         yield mock_datetime