Useful wrappers and helpers

.coveragerc

@@ -1,4 +1,4 @@
 [run]
 relative_files = True
 omit =
-    sekoia_automation/scripts/new_module/template/*
+    sekoia_automation/scripts/new_module/template/*

.github/workflows/ci.yml

@@ -51,12 +51,12 @@ jobs:
 
       - name: Install dependencies
         id: install-dependencies
-        run: poetry install
+        run: poetry install --extras=all
 
       - name: Execute Python tests
         id: execute-tests
         run: |
-          poetry run python -m pytest --junit-xml=junit.xml --cov-report term --cov-report xml:coverage.xml --cov=sekoia_automation tests
+          poetry run python -m pytest --junit-xml=junit.xml --cov-report term --cov-report xml:coverage.xml --cov=sekoia_automation
 
       - name: Upload Test Results
         if: always()

.pre-commit-config.yaml

@@ -1,22 +1,25 @@
 # See https://pre-commit.com for more information
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
-    - repo: https://github.com/charliermarsh/ruff-pre-commit
-        # Ruff version.
-      rev: 'v0.0.285'
+    - repo: local
       hooks:
-          - id: ruff
-            args: [ --fix, --exit-non-zero-on-fix ]
+        - id: black
+          name: Format with Black
+          entry: poetry run black
+          language: system
+          types: [python]
 
-    - repo: https://github.com/ambv/black
-      rev: 23.3.0
-      hooks:
-          - id: black
-            language_version: python3
+        - id: ruff
+          name: Format with Ruff
+          entry: poetry run ruff
+          language: system
+          types: [ python ]
+          args: [ --fix, --exit-non-zero-on-fix, . ]
 
-    - repo: https://github.com/pre-commit/mirrors-mypy
-      rev: 'v1.3.0'  # Use the sha / tag you want to point at
-      hooks:
         - id: mypy
-          args: [--install-types, --non-interactive]
-          exclude: sekoia_automation/scripts/new_module/template/
+          name: Validate types with MyPy
+          entry: poetry run mypy
+          language: system
+          types: [ python ]
+          pass_filenames: false
+          args: [ . ]

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add new wrappers to work with aio libraries such as aiohttp, aiobotocore, etc.
+- New AsyncConnector that contains async implementation of push events
+
 ## [1.4.1] - 2023-09-13
 
 ### Added

pyproject.toml

@@ -44,13 +44,55 @@ PyYAML = "^6.0"
 Jinja2 = "^3.0.3"
 black = "*"  # To format files in cli tools
 prometheus-client = "^0.16.0"
+aiohttp = { version = "^3.8.4", optional = true }
+aiolimiter = { version = "^1.1.0", optional = true }
+aiobotocore = { version = "^2.5.2", optional = true }
+aiofiles = { version = "^23.1.0", optional = true }
+aiocsv = { version = "^1.2.4", optional = true }
+loguru = { version = "^0.7.0", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 unittest-xml-reporting = "^3"
 pylint = "*"
 pytest = "*"
 pytest-cov = "*"
+pytest-asyncio = "*"
+pytest-xdist = "*"
+pytest-env = "*"
 requests-mock = "^1.9"
+faker = "^19.0.0"
+aioresponses = "^0.7.4"
+types-aiofiles = "^23.1.0.4"
+types-requests = "^2.31.0.1"
+types-pyyaml = "^6.0.12.10"
+types-python-slugify = "^8.0.0.2"
+pre-commit = "^3.3.3"
+
+[tool.poetry.extras]
+all = [
+    "aiohttp",
+    "aiolimiter",
+    "aiobotocore",
+    "aiofiles",
+    "aiocsv",
+    "loguru",
+]
+async-aws = [
+    "aiobotocore",
+]
+async-http = [
+    "aiohttp",
+    "aiolimiter",
+    "aiofiles",
+]
+async-files = [
+    "aiofiles",
+    "aiocsv"
+]
+logging = [
+    "loguru"
+]
+
 
 [tool.poetry.group.lint.dependencies]
 ruff = "*"
@@ -63,7 +105,15 @@ force-exclude = "tests/expectations/sample_module/main.py|sekoia_automation/scri
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-addopts = "-ra"
+addopts = '''
+    --asyncio-mode=auto
+    --cache-clear
+    --cov=sekoia_automation
+    --cov-report=html
+    --cov-report=term-missing:skip-covered
+    --cov-fail-under=90
+    -ra
+'''
 testpaths = [
     "tests",
 ]
@@ -72,6 +122,7 @@ testpaths = [
 select = ["A", "ARG", "E", "F", "I", "N", "RUF", "UP", "W"]
 exclude = [
     "tests/expectations/sample_module/main.py",
+    "tests/aio/",
     "sekoia_automation/scripts/new_module/template/"
 ]
 
@@ -82,5 +133,8 @@ exclude = [
 python_version = "3.11"
 ignore_missing_imports = true
 show_column_numbers = true
-exclude = "sekoia_automation/scripts/new_module/template/"
+exclude = [
+    "sekoia_automation/scripts/new_module/template/",
+    "tests/"
+]
 disable_error_code = "annotation-unchecked"

sekoia_automation/aio/__init__.py

@@ -0,0 +1 @@
+"""Package contains all utilities and wrappers for asynchronous mode."""

sekoia_automation/aio/connector.py

@@ -0,0 +1,146 @@
+"""Contains connector with async version."""
+from abc import ABC
+from asyncio import AbstractEventLoop, get_event_loop
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from datetime import datetime
+from urllib.parse import urljoin
+
+from aiohttp import ClientSession
+from aiolimiter import AsyncLimiter
+
+from sekoia_automation.connector import Connector, DefaultConnectorConfiguration
+
+
+class AsyncConnector(Connector, ABC):
+    """Async version of Connector."""
+
+    configuration: DefaultConnectorConfiguration
+
+    _event_loop: AbstractEventLoop
+
+    _session: ClientSession | None = None
+    _rate_limiter: AsyncLimiter | None = None
+
+    def __init__(self, event_loop: AbstractEventLoop | None = None, *args, **kwargs):
+        """
+        Initialize AsyncConnector.
+
+        Optionally accepts event_loop to use, otherwise will use default event loop.
+
+        Args:
+            event_loop: AbstractEventLoop | None
+        """
+        super().__init__(*args, **kwargs)
+
+        self._event_loop = event_loop or get_event_loop()
+
+    @classmethod
+    def set_client_session(cls, session: ClientSession) -> None:
+        """
+        Set client session.
+
+        Args:
+            session: ClientSession
+        """
+        cls._session = session
+
+    @classmethod
+    def set_rate_limiter(cls, rate_limiter: AsyncLimiter) -> None:
+        """
+        Set rate limiter.
+
+        Args:
+            rate_limiter:
+        """
+        cls._rate_limiter = rate_limiter
+
+    @classmethod
+    def get_rate_limiter(cls) -> AsyncLimiter:
+        """
+        Get or initialize rate limiter.
+
+        Returns:
+            AsyncLimiter:
+        """
+        if cls._rate_limiter is None:
+            cls._rate_limiter = AsyncLimiter(1, 1)
+
+        return cls._rate_limiter
+
+    @classmethod
+    @asynccontextmanager
+    async def session(cls) -> AsyncGenerator[ClientSession, None]:  # pragma: no cover
+        """
+        Get or initialize client session if it is not initialized yet.
+
+        Returns:
+            ClientSession:
+        """
+        if cls._session is None:
+            cls._session = ClientSession()
+
+        async with cls.get_rate_limiter():
+            yield cls._session
+
+    async def push_data_to_intakes(
+        self, events: list[str]
+    ) -> list[str]:  # pragma: no cover
+        """
+        Custom method to push events to intakes.
+
+        Args:
+            events: list[str]
+
+        Returns:
+            list[str]:
+        """
+        self._last_events_time = datetime.utcnow()
+        batch_api = urljoin(self.configuration.intake_server, "/batch")
+
+        self.log(f"Push {len(events)} events to intakes")
+
+        result_ids = []
+
+        chunks = self._chunk_events(events, self.configuration.chunk_size)
+
+        async with self.session() as session:
+            for chunk_index, chunk in enumerate(chunks):
+                self.log(
+                    "Start to push chunk {} with data count {} to intakes".format(
+                        chunk_index,
+                        len(chunk),
+                    )
+                )
+
+                request_body = {
+                    "intake_key": self.configuration.intake_key,
+                    "jsons": chunk,
+                }
+
+                for attempt in self._retry():
+                    with attempt:
+                        async with session.post(
+                            batch_api,
+                            headers={"User-Agent": self._connector_user_agent},
+                            json=request_body,
+                        ) as response:
+                            if response.status >= 300:
+                                error = await response.text()
+                                error_message = f"Chunk {chunk_index} error: {error}"
+                                exception = RuntimeError(error_message)
+
+                                self.log(message=error_message, level="error")
+                                self.log_exception(exception)
+
+                                raise exception
+
+                            result = await response.json()
+
+                            self.log(
+                                f"Successfully pushed chunk {chunk_index} to intakes"
+                            )
+
+                            result_ids.extend(result.get("event_ids", []))
+
+        return result_ids

sekoia_automation/aio/helpers/__init__.py

@@ -0,0 +1,5 @@
+"""
+Package contains all utilities and useful helpers.
+
+NOTE!!!: each package inside requires additional libraries to be installed.
+"""

sekoia_automation/aio/helpers/aws/__init__.py

@@ -0,0 +1,7 @@
+"""
+Utilities and wrappers to work with aws.
+
+To use this package you need to install additional libraries:
+
+* aiobotocore (https://github.com/aio-libs/aiobotocore)
+"""