Feature: Connectors workflow

CHANGELOG.md

@@ -7,31 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-## 1.18.0 - 2024-11-26
+## 1.19.0 - 2024-11-28
 
 ### Changed
 
-- Add additional values to log events sent to the API
-- In Generic actions, in case of error use the message from the response if available
-
-## 1.17.2 - 2024-11-06
-
-### Fixed
-
-- Fix callback URL file for account validation
-
-## 1.17.1 - 2024-11-04
-
-### Fixed
-
-- Change the way to handle docker image information when publishing a module
-- Fix the module synchronization script
-
-## 1.17.0 - 2024-11-04
-
-### Added
-
-- Add account validation (beta)
+- Improvements for AsyncConnector.
+- Improvements for Async Http workflow
+- Remove duplicated parts and make the code more uniform for async http workflow
 
 ## 1.16.1 - 2024-10-30
 

pyproject.toml

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

sekoia_automation/aio/connector.py

@@ -1,8 +1,9 @@
 """Contains connector with async version."""
 
 import asyncio
+import time
 from abc import ABC
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Sequence
 from contextlib import asynccontextmanager
 from datetime import datetime
 from pathlib import Path
@@ -12,7 +13,11 @@
 from aiolimiter import AsyncLimiter
 
 from sekoia_automation.aio.helpers import limit_concurrency
-from sekoia_automation.connector import Connector, DefaultConnectorConfiguration
+from sekoia_automation.connector import (
+    Connector,
+    DefaultConnectorConfiguration,
+    EventType,
+)
 from sekoia_automation.module import Module
 
 
@@ -44,24 +49,6 @@
         self.max_concurrency_tasks = kwargs.pop("max_concurrency_tasks", 1000)
         super().__init__(module=module, data_path=data_path, *args, **kwargs)
 
-    def set_client_session(self, session: ClientSession) -> None:
-        """
-        Set client session.
-
-        Args:
-            session: ClientSession
-        """
-        self._session = session
-
-    def set_rate_limiter(self, rate_limiter: AsyncLimiter) -> None:
-        """
-        Set rate limiter.
-
-        Args:
-            rate_limiter:
-        """
-        self._rate_limiter = rate_limiter
-
     def get_rate_limiter(self) -> AsyncLimiter:
         """
         Get or initialize rate limiter.
@@ -137,7 +124,7 @@
             return urljoin(self.intake_url, "batch")
 
     async def push_data_to_intakes(
-        self, events: list[str]
+        self, events: Sequence[EventType]
     ) -> list[str]:  # pragma: no cover
         """
         Custom method to push events to intakes.
@@ -164,13 +151,87 @@
 
         return result_ids
 
-    def stop(self, *args, **kwargs):
-        """
-        Stop the connector
-        """
-        loop = asyncio.get_event_loop()
+    async def async_iterate(
+        self,
+    ) -> AsyncGenerator[tuple[list[EventType], datetime | None], None]:
+        """Iterate over events."""
+        raise NotImplementedError  # To avoid type checking error
+
+    async def async_next_run(self) -> None:
+        processing_start = time.time()
+
+        result_last_event_date: datetime | None = None
+        total_number_of_events = 0
+        async for data in self.async_iterate():  # type: ignore
+            events, last_event_date = data
+            if last_event_date:
+                if (
+                    not result_last_event_date
+                    or last_event_date > result_last_event_date
+                ):
+                    result_last_event_date = last_event_date
+
+            if events:
+                total_number_of_events += len(events)
+                await self.push_data_to_intakes(events)
+
+        processing_end = time.time()
+        processing_time = processing_end - processing_start
+
+        # Metric about processing time
+        self.put_forward_events_duration(
+            intake_key=self.configuration.intake_key,
+            duration=processing_time,
+        )
+
+        # Metric about processing count
+        self.put_forwarded_events(
+            intake_key=self.configuration.intake_key, count=total_number_of_events
+        )
+
+        # Metric about events lag
+        if result_last_event_date:
+            lag = (datetime.utcnow() - result_last_event_date).total_seconds()
+            self.put_events_lag(intake_key=self.configuration.intake_key, lag=lag)
+
+        # Compute the remaining sleeping time.
+        # If greater than 0 and no messages where fetched, pause the connector
+        delta_sleep = (self.frequency or 0) - processing_time
+        if total_number_of_events == 0 and delta_sleep > 0:
+            self.log(message=f"Next batch in the future. Waiting {delta_sleep} seconds")
+
+            await asyncio.sleep(delta_sleep)
+
+    async def on_shutdown(self) -> None:
+        """
+        Called when connector is finishing processing.
+
+        Can be used for some resources cleanup.
+
+        Basically it emits shutdown event.
+        """
+
+    # Put infinite arg only to have testing easier
+    async def async_run(self) -> None:  # pragma: no cover
+        """Runs Connector."""
+        while self.running:
+            try:
+                await self.async_next_run()
+            except Exception as e:
+                self.log_exception(
+                    e,
+                    message=f"Error while running connector {self.connector_name}",
+                )
+
+                if self.frequency:
+                    await asyncio.sleep(self.frequency)
 
         if self._session:
-            loop.run_until_complete(self._session.close())
+            await self._session.close()
 
-        super().stop(*args, **kwargs)
+        await self.on_shutdown()
+
+    def run(self) -> None:  # pragma: no cover
+        """Runs Connector."""
+        loop = asyncio.get_event_loop()
+        loop.run_until_complete(self.async_run())