async_sqs_consumer 1.1.0__tar.gz

async_sqs_consumer-1.1.0/PKG-INFO

@@ -0,0 +1,132 @@
+Metadata-Version: 2.3
+Name: async_sqs_consumer
+Version: 1.1.0
+Summary: 
+Keywords: AWS,AWS,SQS,sqs,consumer,async,worker
+Author: Diego Restrepo
+Author-email: Diego Restrepo <drestrepo@fluidattacks.com>
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: aioboto3>=15.5.0,<16
+Requires-Dist: jsonschema>=4.26.0,<5
+Requires-Dist: pyrate-limiter>=4.0.2,<5
+Requires-Python: >=3.10, <4
+Project-URL: Repository, https://github.com/drestrepom/async_sqs_consumer
+Description-Content-Type: text/markdown
+
+# Async SQS consumer
+
+Python asynchronous (**async** / **await**) worker for consuming messages
+from AWS SQS.
+
+This is a hobby project, if you find the project interesting
+any contribution is welcome.
+
+## Usage
+
+You must create an instance of the worker with the url of the queue.
+
+Aws credentials are taken from environment variables, you must set the
+following environment variables. Or you can provide a Context object with the
+aws credentials `async_sqs_consumer.types.Context`
+
+- `AWS_ACCESS_KEY_ID`
+- `AWS_SECRET_ACCESS_KEY`
+
+Example:
+
+You can get the queue url with the follow aws cli command
+`aws sqs get-queue-url --queue-name xxxxxx`
+
+```python
+# test_worker.py
+
+from async_sqs_consumer.worker import (
+    Worker,
+)
+
+worker = Worker(
+    queue_url="https://sqs.us-east-1.amazonaws.com/xxxxxxx/queue_name"
+)
+
+
+@worker.task("report")
+async def report(text: str) -> None:
+    print(text)
+
+if __name__: "__main__":
+    worker.start()
+```
+
+Now you can initialize the worker `python test_worker.py`
+
+Now you need to post a message for the worker to process
+
+```python
+import json
+import boto3
+import uuid
+
+client = boto3.client("sqs")
+
+client.send_message(
+    QueueUrl="https://sqs.us-east-1.amazonaws.com/xxxxxxx/queue_name",
+    MessageBody=json.dumps(
+        {
+            "task": "report",
+            "id": uuid.uuid4().hex,
+            "args": ["hello world"],
+        }
+    ),
+)
+```
+
+Or you can use aioboto3
+
+```python
+import asyncio
+import json
+import aioboto3
+import uuid
+
+
+async def main() -> None:
+    session = aioboto3.Session()
+    async with session.client("sqs") as client:
+        await client.send_message(
+            QueueUrl="https://sqs.us-east-1.amazonaws.com/xxxxxxx/queue_name",
+            MessageBody=json.dumps(
+                {
+                    "task": "report",
+                    "id": uuid.uuid4().hex,
+                    "args": ["hello world"],
+                }
+            ),
+        )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+To publish the messages they must have the following structure
+
+```json
+{
+    "type": "object",
+    "properties": {
+        "task": {"type": "string"},
+        "id": {"type": "string"},
+        "args": {"type": "array"},
+        "kwargs": {"type": "object"},
+        "retries": {"type": "number"},
+        "eta": {"type": "string"},
+        "expires": {"type": "string"},
+    },
+    "required": ["task", "id"],
+}
+```

async_sqs_consumer-1.1.0/README.md

@@ -0,0 +1,112 @@
+# Async SQS consumer
+
+Python asynchronous (**async** / **await**) worker for consuming messages
+from AWS SQS.
+
+This is a hobby project, if you find the project interesting
+any contribution is welcome.
+
+## Usage
+
+You must create an instance of the worker with the url of the queue.
+
+Aws credentials are taken from environment variables, you must set the
+following environment variables. Or you can provide a Context object with the
+aws credentials `async_sqs_consumer.types.Context`
+
+- `AWS_ACCESS_KEY_ID`
+- `AWS_SECRET_ACCESS_KEY`
+
+Example:
+
+You can get the queue url with the follow aws cli command
+`aws sqs get-queue-url --queue-name xxxxxx`
+
+```python
+# test_worker.py
+
+from async_sqs_consumer.worker import (
+    Worker,
+)
+
+worker = Worker(
+    queue_url="https://sqs.us-east-1.amazonaws.com/xxxxxxx/queue_name"
+)
+
+
+@worker.task("report")
+async def report(text: str) -> None:
+    print(text)
+
+if __name__: "__main__":
+    worker.start()
+```
+
+Now you can initialize the worker `python test_worker.py`
+
+Now you need to post a message for the worker to process
+
+```python
+import json
+import boto3
+import uuid
+
+client = boto3.client("sqs")
+
+client.send_message(
+    QueueUrl="https://sqs.us-east-1.amazonaws.com/xxxxxxx/queue_name",
+    MessageBody=json.dumps(
+        {
+            "task": "report",
+            "id": uuid.uuid4().hex,
+            "args": ["hello world"],
+        }
+    ),
+)
+```
+
+Or you can use aioboto3
+
+```python
+import asyncio
+import json
+import aioboto3
+import uuid
+
+
+async def main() -> None:
+    session = aioboto3.Session()
+    async with session.client("sqs") as client:
+        await client.send_message(
+            QueueUrl="https://sqs.us-east-1.amazonaws.com/xxxxxxx/queue_name",
+            MessageBody=json.dumps(
+                {
+                    "task": "report",
+                    "id": uuid.uuid4().hex,
+                    "args": ["hello world"],
+                }
+            ),
+        )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+To publish the messages they must have the following structure
+
+```json
+{
+    "type": "object",
+    "properties": {
+        "task": {"type": "string"},
+        "id": {"type": "string"},
+        "args": {"type": "array"},
+        "kwargs": {"type": "object"},
+        "retries": {"type": "number"},
+        "eta": {"type": "string"},
+        "expires": {"type": "string"},
+    },
+    "required": ["task", "id"],
+}
+```

async_sqs_consumer-1.1.0/async_sqs_consumer/queue.py

@@ -0,0 +1,221 @@
+from .resources import (
+    get_sqs_client,
+    RESOURCE_OPTIONS_SQS,
+    SESSION,
+)
+from .types import (
+    AwsCredentials,
+)
+from .utils import (
+    TASK_NAME_PREFIX,
+)
+from .utils.retry import (
+    retry,
+)
+from aiohttp.client_exceptions import (
+    ClientConnectorError,
+    ClientPayloadError,
+    ServerDisconnectedError,
+    ServerTimeoutError,
+)
+import asyncio
+from asyncio import (
+    get_event_loop,
+    sleep,
+)
+from mypy_boto3_sqs.client import (
+    SQSClient,
+)
+from botocore.exceptions import (
+    ClientError,
+    ConnectTimeoutError,
+    HTTPClientError,
+    ReadTimeoutError,
+)
+from contextlib import (
+    suppress,
+)
+import logging
+from typing import (
+    Any,
+    Callable,
+    Coroutine,
+)
+
+LOGGER = logging.getLogger(__name__)
+NETWORK_ERRORS = (
+    asyncio.TimeoutError,
+    ClientConnectorError,
+    ClientError,
+    ClientPayloadError,
+    ConnectionResetError,
+    ConnectTimeoutError,
+    HTTPClientError,
+    ReadTimeoutError,
+    ServerDisconnectedError,
+    ServerTimeoutError,
+)
+
+
+def set_max_number_of_messages(max_number_of_messages: int | None) -> int:
+    max_number_of_messages = max_number_of_messages or 10
+    return min(10, max_number_of_messages)
+
+
+@retry(exceptions=NETWORK_ERRORS, tries=3, delay=1)
+async def get_queue_messages(
+    *,
+    queue_url: str,
+    credentials: AwsCredentials | None = None,
+    client: SQSClient | None = None,
+    visibility_timeout: int | None = None,
+    wait_time_seconds: int | None = None,
+    max_number_of_messages: int | None = None,
+) -> list[dict[str, object]]:
+    client = client or await get_sqs_client(credentials)
+    max_number_of_messages = set_max_number_of_messages(max_number_of_messages)
+    try:
+        response = await client.receive_message(  # type: ignore[misc]
+            QueueUrl=queue_url,
+            MaxNumberOfMessages=max_number_of_messages,
+            VisibilityTimeout=visibility_timeout or 60,
+            WaitTimeSeconds=wait_time_seconds or 0,
+        )
+    except Exception as exc:  # pylint: disable=broad-except
+        LOGGER.error(
+            "There was an error with the response",
+            extra={
+                "extra": {
+                    "client": client,
+                    "queue_url": queue_url,
+                    "max_number_of_messages": max_number_of_messages,
+                    "exc": exc,
+                }
+            },
+        )
+
+    return response.get("Messages", [])
+
+
+@retry(exceptions=NETWORK_ERRORS, tries=3, delay=1)
+async def delete_messages(
+    queue_url: str,
+    receipt_handle: dict[str, object],
+    credentials: AwsCredentials | None = None,
+) -> None:
+    try:
+        client = await get_sqs_client(credentials)
+        await client.delete_message(
+            ReceiptHandle=receipt_handle,
+            QueueUrl=queue_url,
+        )
+    except NETWORK_ERRORS as exc:
+        LOGGER.exception(
+            "There was an error deleting messages",
+            extra={
+                "extra": {
+                    "receipt_handle": receipt_handle,
+                    "queue_url": queue_url,
+                    "exc": exc,
+                }
+            },
+        )
+
+
+class Queue:  # pylint: disable=too-many-instance-attributes
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        url: str,
+        priority: int | None = None,
+        authentication: AwsCredentials | None = None,
+        polling_interval: float | None = None,
+        visibility_timeout: int | None = None,
+        max_queue_parallel_messages: int | None = None,
+        enabled: bool | None = None,
+    ) -> None:
+        self.url = url
+        self.priority = priority or 1
+        self.authentication = authentication
+        self.polling_interval = polling_interval or 1.0
+        self.visibility_timeout = visibility_timeout or 60
+        self._max_queue_parallel_messages = max_queue_parallel_messages
+        self._polling = False
+        self.enabled = True if enabled is None else enabled
+
+    async def get_messages(
+        self, sqs_client: object, max_number_of_messages: int | None = None
+    ) -> list[dict[str, object]]:
+        with suppress(asyncio.CancelledError):
+            messages = await get_queue_messages(
+                queue_url=self.url,
+                client=sqs_client,
+                visibility_timeout=self.visibility_timeout,
+                max_number_of_messages=max_number_of_messages,
+            )
+            return messages
+        return []
+
+    def is_task_limit_exceeded(
+        self, max_parallel_messages: int | None
+    ) -> bool:
+        number_of_tasks = len(
+            [
+                task
+                for task in asyncio.all_tasks(get_event_loop())
+                if task.get_name().startswith(TASK_NAME_PREFIX)
+            ]
+        )
+        return number_of_tasks > (
+            self._max_queue_parallel_messages or max_parallel_messages or 1024
+        )
+
+    async def start_polling(
+        self,
+        callback: Callable[
+            [dict[str, object], str], Coroutine[object, object, None]
+        ],
+        queue_alias: str,
+        max_parallel_messages: int | None = None,
+    ) -> None:
+        self._polling = self._polling or True
+        async with SESSION.client(**RESOURCE_OPTIONS_SQS) as sqs_client:
+            while self._polling:
+                if self.is_task_limit_exceeded(max_parallel_messages):
+                    await sleep(self.polling_interval)
+                    continue
+
+                with suppress(asyncio.CancelledError):
+                    messages = await get_queue_messages(
+                        queue_url=self.url,
+                        client=sqs_client,
+                        visibility_timeout=self.visibility_timeout,
+                    )
+                    await asyncio.gather(
+                        *[
+                            callback(message, queue_alias)
+                            for message in messages
+                        ]
+                    )
+                    await sleep(self.polling_interval)
+
+    def stop_polling(self) -> None:
+        self._polling = False
+
+    async def delete_messages(
+        self,
+        receipt_handle: dict[str, Any],
+    ) -> None:
+        async with SESSION.client(**RESOURCE_OPTIONS_SQS) as sqs_client:
+            await sqs_client.delete_message(
+                ReceiptHandle=receipt_handle,
+                QueueUrl=self.url,
+            )
+        LOGGER.info(
+            "Deleted messages parameters",
+            extra={
+                "extra": {
+                    "sqs_client": self.url,
+                    "receipt_handle": receipt_handle,
+                }
+            },
+        )

async_sqs_consumer-1.1.0/async_sqs_consumer/resources.py

@@ -0,0 +1,72 @@
+from .types import (
+    AwsCredentials,
+)
+import aioboto3  # type: ignore[import-untyped]
+from aiobotocore.config import (  # type: ignore[import-untyped]
+    AioConfig,
+)
+from contextlib import (
+    AsyncExitStack,
+)
+from typing import (
+    Any,
+)
+
+RESOURCE_OPTIONS_SQS = {
+    "aws_access_key_id": (None),
+    "aws_secret_access_key": (None),
+    "config": AioConfig(
+        # The time in seconds till a timeout exception is thrown when
+        # attempting to make a connection. [60]
+        connect_timeout=15,
+        # Maximum amount of simultaneously opened connections. [10]
+        # https://docs.aiohttp.org/en/stable/client_advanced.html#limiting-connection-pool-size
+        max_pool_connections=2000,
+        # The time in seconds till a timeout exception is thrown when
+        # attempting to read from a connection. [60]
+        read_timeout=30,
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html
+        retries={"max_attempts": 10, "mode": "standard"},
+        # Signature version for signing URLs
+        # https://boto3.amazonaws.com/v1/documentation/api/1.9.42/guide/s3.html#generating-presigned-urls
+        signature_version="s3v4",
+    ),
+    "endpoint_url": None,
+    "region_name": "us-east-1",
+    "service_name": "sqs",
+    "use_ssl": True,
+    "verify": True,
+}
+
+SESSION = aioboto3.Session()
+CONTEXT_STACK_SQS = None
+RESOURCE_SQS = None
+
+
+async def sqs_startup(credentials: AwsCredentials | None = None) -> None:
+    # pylint: disable=global-statement
+    global CONTEXT_STACK_SQS, RESOURCE_SQS
+
+    CONTEXT_STACK_SQS = AsyncExitStack()
+    if credentials:
+        RESOURCE_OPTIONS_SQS["aws_access_key_id"] = credentials.access_key_id
+        RESOURCE_OPTIONS_SQS[
+            "aws_secret_access_key"
+        ] = credentials.secret_access_key
+        RESOURCE_OPTIONS_SQS["aws_session_token"] = credentials.session_token
+
+    RESOURCE_SQS = await CONTEXT_STACK_SQS.enter_async_context(
+        SESSION.client(**RESOURCE_OPTIONS_SQS)
+    )
+
+
+async def sqs_shutdown() -> None:
+    if CONTEXT_STACK_SQS:
+        await CONTEXT_STACK_SQS.aclose()
+
+
+async def get_sqs_client(credentials: AwsCredentials | None = None) -> Any:
+    if RESOURCE_SQS is None:
+        await sqs_startup(credentials)
+
+    return RESOURCE_SQS

async_sqs_consumer-1.1.0/async_sqs_consumer/types.py

@@ -0,0 +1,56 @@
+from enum import (
+    Enum,
+)
+from typing import (
+    NamedTuple,
+)
+
+MESSAGE_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "task": {"type": "string"},
+        "id": {"type": "string"},
+        "args": {"type": "array"},
+        "kwargs": {"type": "object"},
+        "retries": {"type": "number"},
+        "eta": {"type": "string"},
+        "expires": {"type": "string"},
+    },
+    "required": ["task", "id"],
+}
+SNS_NOTIFICATION_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "Type": {"type": "string"},
+        "MessageId": {"type": "string"},
+        "TopicArn": {"type": "string"},
+        "Message": {"type": "string"},
+        "Timestamp": {"type": "string"},
+        "SignatureVersion": {"type": "string"},
+        "Signature": {"type": "string"},
+        "SigningCertURL": {"type": "string"},
+        "UnsubscribeURL": {"type": "string"},
+    },
+    "required": ["Type", "Message", "Timestamp"],
+}
+
+
+class AwsCredentials(NamedTuple):
+    access_key_id: str
+    secret_access_key: str
+    session_token: str | None = None
+
+
+class MessageType(Enum):
+    CUSTOM = "CUSTOM"
+    SNS_NOTIFICATION = "SNS_NOTIFICATION"
+
+
+class SqsOptions(NamedTuple):
+    visibility_timeout: int | None = None
+
+
+class Context(NamedTuple):
+    queue_url: str
+    aws_credentials: AwsCredentials | None = None
+    sqs: SqsOptions | None = None

async_sqs_consumer-1.1.0/async_sqs_consumer/utils/__init__.py

@@ -0,0 +1,47 @@
+from ..types import (
+    MESSAGE_SCHEMA,
+    MessageType,
+    SNS_NOTIFICATION_SCHEMA,
+)
+from contextlib import (
+    suppress,
+)
+from jsonschema import (
+    validate,
+)
+from jsonschema.exceptions import (
+    ValidationError,
+)
+
+TASK_NAME_PREFIX = "task_queue"
+
+
+class FailedValidation(Exception):
+    pass
+
+
+def validate_message(
+    message: dict[str, object], message_type: MessageType
+) -> bool:
+    with suppress(ValidationError):
+        if message_type == MessageType.CUSTOM:
+            validate(message, MESSAGE_SCHEMA)
+        else:
+            validate(message, SNS_NOTIFICATION_SCHEMA)
+        return True
+    return False
+
+
+def inverse_proportional_distribution(
+    amount: int, priorities: list[int]
+) -> list[float]:
+    inv_priorities = [1 / p for p in priorities]
+    total_inv_priority = sum(inv_priorities)
+    inv_proportions = [inv_p / total_inv_priority for inv_p in inv_priorities]
+    amounts = [amount * p for p in inv_proportions]
+    total_amounts = sum(amounts)
+    amounts = [
+        int(round(a + (amount - total_amounts) / len(amounts)))
+        for a in amounts
+    ]
+    return amounts

async_sqs_consumer-1.1.0/async_sqs_consumer/utils/retry.py

@@ -0,0 +1,152 @@
+# pylint: disable=broad-except
+# pylint: disable=too-many-arguments
+import asyncio
+import functools
+from functools import (
+    partial,
+)
+import logging
+import random
+import traceback
+from typing import (
+    Any,
+    Callable,
+    Literal,
+    Type,
+)
+
+logging_logger = logging.getLogger(__name__)
+
+
+def _log_exception(
+    func: partial,
+    exc: Exception,
+    logger: logging.Logger | None = None,
+    log_traceback: bool = False,
+    _delay: float | int | None = None,
+) -> None:
+    if logger is not None:
+        try:
+            func_qualname = func.func.__qualname__
+        except AttributeError:
+            func_qualname = str(func.func)
+        logger.warning(
+            "%s: %s in %s.%s, retrying in %s seconds...",
+            exc.__class__.__qualname__,
+            exc,
+            func.func.__module__,
+            func_qualname,
+            _delay,
+        )
+        if log_traceback:
+            logger.warning(traceback.format_exc())
+
+
+def _increase_delay(
+    _delay: float,
+    backoff: float | None = None,
+    jitter: Literal[1] | Literal[0] = 0,
+    max_delay: float | None = None,
+) -> float:
+    _delay *= backoff or 1
+
+    if isinstance(jitter, tuple):
+        _delay += random.uniform(*jitter)
+    else:
+        _delay += jitter
+
+    if max_delay is not None:
+        _delay = min(_delay, max_delay)
+
+    return _delay
+
+
+async def __retry_internal(
+    func: partial,
+    tries: int,
+    exceptions: Type[Exception] | tuple[type[Exception], ...] | None = None,
+    delay: float | None = None,
+    max_delay: float | None = None,
+    backoff: float | None = None,
+    jitter: Literal[1] | Literal[0] = 0,
+    logger: logging.Logger | None = None,
+    log_traceback: bool = False,
+    on_exception: Callable[[Exception], bool] | None = None,
+) -> None:
+    exceptions = exceptions or Exception
+    _tries, _delay = tries, (delay or 0.0)
+    logger = logger or logging_logger
+    while _tries >= 0:
+        try:
+            return await func()
+        except exceptions as exc:
+            if on_exception is not None and on_exception(exc):
+                break
+
+            _tries -= 1
+            _log_exception(func, exc, logger, log_traceback, _delay)
+
+            await asyncio.sleep(_delay)
+            _delay = _increase_delay(_delay, backoff, jitter, max_delay)
+
+
+def retry(  # pylint: disable=too-many-arguments
+    tries: int,
+    exceptions: Type[Exception] | tuple[type[Exception], ...] = Exception,
+    delay: float | None = None,
+    max_delay: float | None = None,
+    backoff: float | None = None,
+    jitter: Literal[1] | Literal[0] = 0,
+    logger: logging.Logger | None = None,
+    log_traceback: bool = False,
+    on_exception: Callable[[Exception], bool] | None = None,
+) -> Callable[[Callable], Callable]:
+    def decorator(func: Callable[..., str]) -> Any:
+        @functools.wraps(func)
+        async def wrapper(
+            *fargs: list[str], **fkwargs: dict[str, str]
+        ) -> None:
+            args = fargs if fargs else []
+            kwargs = fkwargs if fkwargs else {}
+            return await __retry_internal(
+                partial(func, *args, **kwargs),
+                tries,
+                exceptions,
+                delay,
+                max_delay,
+                backoff,
+                jitter,
+                logger,
+                log_traceback,
+                on_exception,
+            )
+
+        return wrapper
+
+    return decorator
+
+
+async def retry_call(  # pylint: disable=too-many-arguments
+    func: partial | Callable[..., object],
+    tries: int,
+    exceptions: Type[Exception] | tuple[type[Exception], ...] = Exception,
+    delay: float | None = None,
+    max_delay: float | None = None,
+    backoff: float | None = None,
+    jitter: Literal[1] | Literal[0] = 0,
+    logger: logging.Logger | None = None,
+    fargs: tuple[object, ...] | None = None,
+    fkwargs: dict[str, object] | None = None,
+) -> None:
+    args = fargs if fargs else tuple()
+    kwargs = fkwargs if fkwargs else {}
+    return await __retry_internal(
+        partial(func, *args, **kwargs),
+        tries,
+        exceptions,
+        delay,
+        max_delay,
+        backoff,
+        jitter,
+        logger,
+    )

async_sqs_consumer-1.1.0/async_sqs_consumer/worker.py

@@ -0,0 +1,377 @@
+from .queue import (
+    Queue,
+)
+from .resources import (
+    RESOURCE_OPTIONS_SQS,
+    SESSION,
+)
+from .types import (
+    MessageType,
+)
+from .utils import (
+    inverse_proportional_distribution,
+    TASK_NAME_PREFIX,
+    validate_message,
+)
+from .utils.retry import (
+    retry_call,
+)
+import asyncio
+from contextlib import (
+    suppress,
+)
+from datetime import (
+    datetime,
+    timezone,
+)
+from functools import (
+    partial,
+)
+import json
+import logging
+import math
+import signal
+import sys
+from typing import (
+    Any,
+    Callable,
+    Literal,
+)
+from uuid import (
+    uuid4,
+)
+
+logging.basicConfig(level=logging.INFO)
+
+LOGGER = logging.getLogger(__name__)
+
+
+class Worker:
+    def __init__(
+        self, queues: dict[str, Queue], max_workers: int | None = None
+    ) -> None:
+        self.queues = queues
+        self.running = False
+        self.handlers_by_queue: dict[
+            str, dict[str, Callable[..., object]]
+        ] = {}
+        self._events: dict[str, list[Callable[[], None]]] = {}
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._tasks: list[asyncio.Task] = []
+        self._max_workers = max_workers or 1024
+
+    def _sigint_handler(self, *_args: object) -> None:
+        sys.stdout.write("\b\b\r")
+        sys.stdout.flush()
+        logging.getLogger("system").warning(
+            "Received <ctrl+c> interrupt [SIGINT]"
+        )
+        self._stop_worker()
+
+    def _sigterm_handler(self, *_args: object) -> None:
+        logging.getLogger("system").warning(
+            "Received termination signal [SIGTERM]"
+        )
+        self._stop_worker()
+
+    def _stop_worker(self, *_: object) -> None:
+        if self._loop:
+            LOGGER.debug("Stop worker")
+            self.running = False
+            for queue in self.queues.values():
+                queue.stop_polling()
+            for event in self._events.get("shutdown", []):
+                self._loop.run_until_complete(event())  # type: ignore
+
+            map(lambda task: task.cancel(), self._tasks)
+
+            # make sure all tasks are finished before finishing the
+            # entire process
+            self._loop.stop()
+
+    def task(
+        self, name: str, queue_name: str | None = None
+    ) -> Callable[..., None]:
+        """A decorator to add a task that could be handled by the worker.
+
+        Args:
+            name (str): tasks received with this name will be
+            processed by the decorated function
+        Returns:
+            Callable[[Callable[[], None]], None]: _description_
+        """
+
+        queue_name = queue_name or "default"
+
+        def decorator(func: Callable[[], object]) -> None:
+            queue_ = queue_name or "default"
+            if queue_ not in self.handlers_by_queue:
+                self.handlers_by_queue[queue_] = {name: func}
+            else:
+                self.handlers_by_queue[queue_][name] = func
+
+        return decorator
+
+    def on_event(
+        self, event_name: Literal["startup", "shutdown"]
+    ) -> Callable[[Callable[[], None]], None]:
+        """Yo can define event handlers that need to be executed before
+        the application startups, or when the application is shutting down.
+        The valid events are `startup`, `shutdown`.
+
+        Args:
+            event_name (Literal[&quot;startup&quot;, &quot;shutdown&quot;]):
+
+        Returns:
+            Callable[[Callable[[], None]], None]:
+        """
+
+        def decorator(func: Callable[[], None]) -> None:
+            if event_name not in self._events:
+                self._events[event_name] = [func]
+            else:
+                self._events[event_name] = [*self._events[event_name], func]
+
+        return decorator
+
+    def start(
+        self, event_loop: asyncio.AbstractEventLoop | None = None
+    ) -> None:
+        """Start the workert
+
+        Args:
+            event_loop (asyncio.AbstractEventLoop  | None, optional):
+            if an eventloop is already running you can run the worker in it.
+            Defaults to None.
+        """
+        self.running = True
+        loop = event_loop or asyncio.get_event_loop()
+        self._loop = loop
+
+        if self._loop and self._loop.is_closed():
+            self._loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(self._loop)
+
+        for signame in ("SIGINT", "SIGTERM"):
+            self._loop.add_signal_handler(
+                getattr(signal, signame), partial(self._stop_worker, self)
+            )
+
+        signal.siginterrupt(signal.SIGTERM, False)
+        signal.siginterrupt(signal.SIGUSR1, False)
+        signal.signal(signal.SIGINT, partial(self._sigint_handler, self))
+        signal.signal(signal.SIGTERM, partial(self._sigterm_handler, self))
+
+        for event in self._events.get("startup", []):
+            self._loop.run_until_complete(event())  # type: ignore
+        loop.create_task(self._poll_messages())
+        self._loop.run_forever()
+
+    def _finish_message(
+        self,
+        _future: asyncio.Future | None,
+        *,
+        handler_name: str,
+        task_id: str,
+        queue_alias: str,
+        start_time: float,
+        receipt_handle: object | None = None,
+    ) -> None:
+        start_date = datetime.fromtimestamp(start_time, tz=timezone.utc)
+        end_date = datetime.now(tz=timezone.utc)
+        total_seconds = end_date - start_date
+        if _future and _future.exception() is not None:
+            LOGGER.warning(
+                "Failed to execute task %s[%s]", handler_name, task_id
+            )
+            LOGGER.exception(
+                "Following exception",
+                exc_info=_future.exception(),
+            )
+        else:
+            LOGGER.info(
+                "Task %s[%s] succeeded in %ss: None",
+                handler_name,
+                task_id,
+                total_seconds.seconds,
+            )
+            if receipt_handle:
+                self._delete_message(receipt_handle, queue_alias)
+
+    def _delete_message(
+        self,
+        receipt_handle: object,
+        queue_alias: str,
+    ) -> None:
+        asyncio.ensure_future(
+            self.queues[queue_alias].delete_messages(
+                receipt_handle  # type: ignore[arg-type]
+            ),
+            loop=self._loop,
+        )
+
+    def _available_workers_to_consume(self) -> int:
+        current_tasks = [
+            task
+            for task in asyncio.all_tasks(asyncio.get_event_loop())
+            if task.get_name().startswith(TASK_NAME_PREFIX)
+        ]
+        return self._max_workers - len(current_tasks)
+
+    def _get_tasks_by_queue(
+        self, queue_alias: str
+    ) -> tuple[asyncio.Task, ...]:
+        return tuple(
+            task
+            for task in asyncio.all_tasks(asyncio.get_event_loop())
+            if task.get_name().startswith(f"{TASK_NAME_PREFIX}_{queue_alias}")
+        )
+
+    async def queue_message_to_worker(
+        self,
+        *,
+        body: dict[str, Any],
+        message_type: MessageType,
+        queue_alias: str,
+        receipt_handler: str,
+    ) -> bool:
+        task_id = body.get("id") or uuid4().hex
+        task_name = body.get("task") or message_type.value.lower()
+        LOGGER.info("Task %s[%s] received", task_name, task_id)
+
+        if not self._loop:
+            return False
+
+        handler = self.handlers_by_queue.get(queue_alias, {}).get(task_name)
+        if not handler:
+            LOGGER.warning(
+                "The task %s has not a candidate to be proceeded",
+                task_name,
+            )
+            return False
+
+        task = self._loop.create_task(
+            retry_call(
+                handler,
+                fargs=tuple(body.get("args", [])),
+                fkwargs=(
+                    body
+                    if message_type == MessageType.SNS_NOTIFICATION
+                    else body.get("kwargs", {})
+                ),
+                tries=body.get("retries", 1),
+            ),
+            name=(
+                f"{TASK_NAME_PREFIX}_{queue_alias}"
+                f"_{task_name}_{uuid4().hex[:8]}"
+            ),
+        )
+        task.add_done_callback(
+            partial(
+                self._finish_message,
+                receipt_handle=receipt_handler,
+                queue_alias=queue_alias,
+                start_time=datetime.now().timestamp(),
+                handler_name=task_name,
+                task_id=task_id,
+            )
+        )
+        return True
+
+    async def _process_queue_message(
+        self, message_content: dict[str, Any], queue_alias: str
+    ) -> bool:
+        try:
+            body = json.loads(message_content["Body"])
+        except json.JSONDecodeError:
+            return False
+
+        message_type = (
+            MessageType.SNS_NOTIFICATION
+            if body.get("Type") == "Notification"
+            else MessageType.CUSTOM
+        )
+        if not validate_message(body, message_type):
+            LOGGER.error("Failed to validate message")
+            return False
+
+        await self.queue_message_to_worker(
+            body=body,
+            message_type=message_type,
+            queue_alias=queue_alias,
+            receipt_handler=message_content["ReceiptHandle"],
+        )
+
+        return True
+
+    async def _consumer_callback(
+        self, message: dict[str, Any], queue_alias: str
+    ) -> None:
+        with suppress(asyncio.CancelledError):
+            success = await self._process_queue_message(message, queue_alias)
+            if not success:
+                await self.queues[queue_alias].delete_messages(
+                    message["ReceiptHandle"]
+                )
+
+    async def _queue_n_messages(
+        self,
+        *,
+        queue: Queue,
+        queue_alias: str,
+        sqs_client: object,
+        required_messages: int,
+    ) -> None:
+        for _ in range(math.ceil(required_messages / 10)):
+            messages = await queue.get_messages(
+                sqs_client,
+                max_number_of_messages=(required_messages),
+            )
+
+            if len(messages) == 0:
+                break
+
+            await asyncio.gather(
+                *[
+                    self._consumer_callback(message, queue_alias)
+                    for message in messages
+                ]
+            )
+
+    async def _poll_messages(self) -> None:
+        queues = list(sorted(self.queues.items(), key=lambda x: x[1].priority))
+        queues = [item for item in queues if item[1].enabled]
+        async with SESSION.client(**RESOURCE_OPTIONS_SQS) as sqs_client:
+            while True:
+                if self._available_workers_to_consume() < 1:
+                    await asyncio.sleep(1)
+                    continue
+                for index, (queue_alias, queue) in enumerate(queues):
+                    required_messages = math.ceil(
+                        inverse_proportional_distribution(
+                            self._available_workers_to_consume(),
+                            [q[1].priority for q in queues[index:]],
+                        )[0]
+                    )
+
+                    # it is necessary to request the messages several times
+                    # because sqs only allows to obtain 10 messages at the
+                    # same time
+                    await self._queue_n_messages(
+                        queue=queue,
+                        queue_alias=queue_alias,
+                        sqs_client=sqs_client,
+                        required_messages=required_messages,
+                    )
+
+                    LOGGER.info(
+                        "Queue messages info",
+                        extra={
+                            "extra": {
+                                "queue": queue,
+                                "queue_alias": queue_alias,
+                                "sqs_client": sqs_client,
+                                "required_messages": required_messages,
+                            }
+                        },
+                    )

async_sqs_consumer-1.1.0/pyproject.toml

@@ -0,0 +1,51 @@
+[project]
+name = "async_sqs_consumer"
+version = "1.1.0"
+description = ""
+authors = [{ name = "Diego Restrepo", email = "drestrepo@fluidattacks.com" }]
+requires-python = ">=3.10,<4"
+readme = "README.md"
+keywords = [
+    "AWS",
+    "AWS",
+    "SQS",
+    "sqs",
+    "consumer",
+    "async",
+    "worker",
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+]
+dependencies = [
+    "aioboto3>=15.5.0,<16",
+    "jsonschema>=4.26.0,<5",
+    "pyrate-limiter>=4.0.2,<5",
+]
+
+[project.urls]
+Repository = "https://github.com/drestrepom/async_sqs_consumer"
+
+[dependency-groups]
+dev = [
+    "prospector>=1.7.7,<2",
+    "black>=22.8.0,<23",
+    "isort>=5.10.1,<6",
+    "mypy>=1.19.1",
+    "types-jsonschema>=4.21.0",
+    "boto3-stubs[essential]>=1.34.0",
+]
+
+[tool.uv]
+
+[tool.uv.build-backend]
+module-root = ""
+
+[build-system]
+requires = ["uv_build>=0.9.26,<0.10.0"]
+build-backend = "uv_build"