ergon-framework-python 0.1.0__py3-none-any.whl

ergon/connector/rabbitmq/service.py

@@ -0,0 +1,199 @@
+import functools
+import json
+import logging
+import ssl as ssl_module
+import time
+from collections import deque
+from typing import Any, Dict, List, Optional
+
+import pika
+
+from .helper import headers_generator
+from .models import RabbitmqClient, RabbitmqProducerMessage
+
+logger = logging.getLogger(__name__)
+
+
+class RabbitMQService:
+    def __init__(self, client: RabbitmqClient) -> None:
+        self.client = client
+
+        self._connection: Optional[pika.BlockingConnection] = None
+        self._channel: Any = None
+        self._connect()
+
+    # ---------- Connection / Channel ----------
+
+    def _connect(self) -> None:
+        """Ensure an open connection and channel exist."""
+        if self._connection and self._connection.is_open:
+            return
+
+        ssl_options = None
+        if self.client.ssl_enabled:
+            ctx = ssl_module.create_default_context()
+            if self.client.ssl_ca_certs:
+                ctx.load_verify_locations(self.client.ssl_ca_certs)
+            ssl_options = pika.SSLOptions(ctx)
+
+        params = pika.ConnectionParameters(
+            host=self.client.host,
+            port=self.client.port,
+            virtual_host=self.client.virtual_host or "/",
+            credentials=pika.PlainCredentials(self.client.username, self.client.password),
+            connection_attempts=self.client.connection_attempts,
+            socket_timeout=self.client.socket_timeout,
+            heartbeat=int(self.client.heartbeat) if self.client.heartbeat is not None else None,
+            blocked_connection_timeout=self.client.blocked_connection_timeout,
+            ssl_options=ssl_options,
+        )
+        self._connection = pika.BlockingConnection(params)
+        self._channel = self._connection.channel()
+
+        self._channel.queue_declare(queue=self.client.queue_name, durable=True)
+        self._channel.basic_qos(prefetch_count=self.client.prefetch_count)
+
+    def close(self) -> None:
+        """
+        Fecha conexão/canal, se estiverem abertos.
+        """
+        try:
+            if self._channel and self._channel.is_open:
+                self._channel.close()
+        finally:
+            if self._connection and self._connection.is_open:
+                self._connection.close()
+
+    # ---------- Consumo ----------
+
+    def consume(
+        self,
+        queue_name,
+        auto_ack,
+        batch_size=1,
+    ) -> List[Dict[str, Any]]:
+        """
+        Gera mensagens da fila como dicionários.
+
+        Cada item gerado pode conter:
+        - payload: corpo da mensagem (json decodado, se possível)
+        - routing_key
+        - delivery_tag
+        - headers
+        """
+        assert self._connection is not None
+        assert self._channel is not None
+
+        q = queue_name or self.client.queue_name
+
+        buffer = deque()
+
+        # Callback interno que vai encher o buffer
+        def _internal_callback(ch, method, properties, body):
+            try:
+                try:
+                    payload = json.loads(body.decode("utf-8")) if body else None
+                except Exception:
+                    payload = body
+
+                if len(buffer) >= batch_size:
+                    # Já tenho batch, só requeue ou ignora
+                    if not auto_ack:
+                        ch.basic_nack(method.delivery_tag, requeue=True)
+                    return
+
+                buffer.append(
+                    {
+                        "data": payload,
+                        "routing_key": method.routing_key,
+                        "delivery_tag": method.delivery_tag,
+                        "headers": getattr(properties, "headers", {}) or {},
+                    }
+                )
+
+                if auto_ack:
+                    ch.basic_ack(method.delivery_tag)
+
+                # Cancela consumo após capturar o batch
+                if len(buffer) == batch_size:
+                    self._channel.basic_cancel(consumer_tag)
+            except Exception:
+                if not auto_ack:
+                    ch.basic_nack(method.delivery_tag, requeue=True)
+
+        # Registra o callback
+        consumer_tag = self._channel.basic_consume(
+            queue=q,
+            on_message_callback=_internal_callback,
+            auto_ack=auto_ack,
+        )
+
+        # Loop para acumular mensagens
+        start = time.time()
+        timeout = 2.0
+
+        try:
+            while len(buffer) < batch_size:
+                self._connection.process_data_events(time_limit=0.1)  # type: ignore[union-attr, arg-type]
+                # se houver menos mensagem que o tamanho do batch sai do loop por timeout
+                if time.time() - start > timeout:
+                    break
+        finally:
+            self._channel.basic_cancel(consumer_tag)
+
+        return list(buffer)
+
+    def ack_msg(self, delivery_tag) -> None:
+        assert self._channel is not None
+        try:
+            self._channel.basic_ack(delivery_tag=delivery_tag)
+            logger.info(f"sucesso ao marcar como lida. tag {delivery_tag}")
+        except Exception as e:
+            logger.error("Falha ao ackar a mensagem")
+            logger.error(f"{str(e)}")
+            self._channel.basic_nack(delivery_tag, requeue=True)
+
+    def confirm_message_received(self, delivery_tag: int) -> None:
+        if not isinstance(delivery_tag, int):
+            try:
+                delivery_tag = int(delivery_tag)
+            except ValueError:
+                raise ValueError("Delivery tag must be an integer.")
+
+        assert self._connection is not None
+        logger.info(f"marcando mensagem como recebida. tag {delivery_tag}")
+        ack = functools.partial(self.ack_msg, delivery_tag=delivery_tag)
+        self._connection.add_callback_threadsafe(ack)
+
+    # ---------- Publish ----------
+
+    def publish(self, message: RabbitmqProducerMessage, exchange: str = "") -> None:
+        if not (self._channel and self._channel.is_open):
+            logger.warning("Channel closed, attempting reconnect before publish")
+            self._connect()
+            if not (self._channel and self._channel.is_open):
+                raise ConnectionError("Unable to publish: RabbitMQ channel is closed after reconnect attempt")
+
+        rk = message.payload.get("queue_name") or self.client.queue_name  # type: ignore[attr-defined]
+
+        raw = message.payload.get("body")  # type: ignore[attr-defined]
+        if not isinstance(raw, (str, bytes)):
+            body = json.dumps(raw)
+        elif isinstance(raw, str):
+            body = raw.encode("utf-8")
+        else:
+            body = json.dumps(raw, ensure_ascii=False).encode("utf-8")
+
+        properties = pika.BasicProperties(
+            headers=headers_generator(id=message.id, source=message.payload.get("source")),  # type: ignore[attr-defined]
+            delivery_mode=message.delivery_mode,
+            content_type=message.content_type,
+        )
+
+        self._channel.basic_publish(
+            exchange=exchange,
+            routing_key=rk,
+            body=body,
+            properties=properties,
+        )
+        logger.info("Message published successfully")

ergon/connector/sqs/__init__.py

@@ -0,0 +1,15 @@
+from .async_connector import AsyncSQSConnector
+from .async_service import AsyncSQSService
+from .connector import SQSConnector
+from .models import SQSClient, SQSConsumerConfig, SQSProducerConfig
+from .service import SQSService
+
+__all__ = [
+    "AsyncSQSConnector",
+    "AsyncSQSService",
+    "SQSConnector",
+    "SQSService",
+    "SQSClient",
+    "SQSConsumerConfig",
+    "SQSProducerConfig",
+]

ergon/connector/sqs/async_connector.py

@@ -0,0 +1,120 @@
+import json
+import logging
+from typing import Any, Dict, List, Optional
+
+from ..connector import AsyncConnector
+from ..transaction import Transaction
+from .async_service import AsyncSQSService
+from .models import SQSClient, SQSConsumerConfig, SQSProducerConfig
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncSQSConnector(AsyncConnector):
+    service: AsyncSQSService
+
+    def __init__(
+        self,
+        client: SQSClient,
+        consumer_config: Optional[SQSConsumerConfig] = None,
+        producer_config: Optional[SQSProducerConfig] = None,
+    ) -> None:
+        self.service = AsyncSQSService(client)
+        self._consumer_config = consumer_config or SQSConsumerConfig()
+        self._producer_config = producer_config or SQSProducerConfig()
+
+    async def fetch_transactions_async(
+        self,
+        batch_size: int = 1,
+        queue_url: Optional[str] = None,
+        wait_time_seconds: Optional[int] = None,
+        visibility_timeout: Optional[int] = None,
+        attribute_names: Optional[List[str]] = None,
+        message_attribute_names: Optional[List[str]] = None,
+        *args,
+        **kwargs,
+    ) -> List[Transaction]:
+        cfg = self._consumer_config
+
+        resolved_queue_url = queue_url or cfg.queue_url
+        resolved_wait_time = wait_time_seconds if wait_time_seconds is not None else cfg.wait_time_seconds
+        resolved_visibility = visibility_timeout if visibility_timeout is not None else cfg.visibility_timeout
+        resolved_attr_names = attribute_names or cfg.attribute_names
+        resolved_msg_attr_names = message_attribute_names or cfg.message_attribute_names
+
+        messages = await self.service.receive_messages(
+            queue_url=resolved_queue_url,
+            max_number_of_messages=batch_size,
+            wait_time_seconds=resolved_wait_time,
+            visibility_timeout=resolved_visibility,
+            attribute_names=resolved_attr_names,
+            message_attribute_names=resolved_msg_attr_names,
+        )
+
+        if not messages:
+            return []
+
+        effective_queue_url = resolved_queue_url or self.service.client.queue_url
+
+        return [
+            Transaction(
+                id=msg["MessageId"],
+                payload=msg,
+                metadata={
+                    "receipt_handle": msg["ReceiptHandle"],
+                    "queue_url": effective_queue_url,
+                },
+            )
+            for msg in messages
+        ]
+
+    async def dispatch_transactions_async(
+        self,
+        transactions: List[Transaction],
+        queue_url: Optional[str] = None,
+        delay_seconds: Optional[int] = None,
+        message_group_id: Optional[str] = None,
+        message_deduplication_id: Optional[str] = None,
+        message_attributes: Optional[Dict[str, Dict]] = None,
+        *args,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        cfg = self._producer_config
+
+        resolved_queue_url = queue_url or cfg.queue_url
+        resolved_delay = delay_seconds if delay_seconds is not None else cfg.delay_seconds
+        resolved_group_id = message_group_id or cfg.message_group_id
+        resolved_dedup_id = message_deduplication_id or cfg.message_deduplication_id
+        resolved_msg_attrs = message_attributes or cfg.message_attributes
+
+        entries: List[Dict[str, Any]] = []
+        for i, txn in enumerate(transactions):
+            body = txn.payload
+            if not isinstance(body, str):
+                body = json.dumps(body)
+
+            entry: Dict[str, Any] = {
+                "Id": str(i),
+                "MessageBody": body,
+                "DelaySeconds": resolved_delay,
+            }
+            if resolved_msg_attrs:
+                entry["MessageAttributes"] = resolved_msg_attrs
+            if resolved_group_id:
+                entry["MessageGroupId"] = resolved_group_id
+
+            txn_dedup_id = txn.metadata.get("message_deduplication_id") if txn.metadata else None
+            dedup_id = txn_dedup_id or resolved_dedup_id
+            if dedup_id:
+                entry["MessageDeduplicationId"] = dedup_id
+
+            entries.append(entry)
+
+        return await self.service.send_message_batch(entries=entries, queue_url=resolved_queue_url)
+
+    async def get_transactions_count_async(self, queue_url: Optional[str] = None, *args, **kwargs) -> int:
+        attrs = await self.service.get_queue_attributes(
+            attribute_names=["ApproximateNumberOfMessages"],
+            queue_url=queue_url,
+        )
+        return int(attrs.get("ApproximateNumberOfMessages", 0))

ergon/connector/sqs/async_service.py

@@ -0,0 +1,246 @@
+import json
+import logging
+from typing import Any, Dict, List, Optional
+
+from aiobotocore.session import get_session
+
+from .models import SQSClient
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncSQSService:
+    _sqs: Any
+    _sqs_ctx: Any
+
+    def __init__(self, client: SQSClient) -> None:
+        self.client = client
+
+        self._session = get_session()
+        self._client_kwargs: Dict[str, Any] = {"region_name": client.region_name}
+        if client.aws_access_key_id:
+            self._client_kwargs["aws_access_key_id"] = client.aws_access_key_id
+        if client.aws_secret_access_key:
+            self._client_kwargs["aws_secret_access_key"] = client.aws_secret_access_key
+        if client.aws_session_token:
+            self._client_kwargs["aws_session_token"] = client.aws_session_token
+        if client.endpoint_url:
+            self._client_kwargs["endpoint_url"] = client.endpoint_url
+
+        self._sqs_ctx = None
+        self._sqs = None
+
+    async def _get_client(self) -> Any:
+        if self._sqs is None:
+            self._sqs_ctx = self._session.create_client("sqs", **self._client_kwargs)
+            self._sqs = await self._sqs_ctx.__aenter__()
+        return self._sqs
+
+    async def close(self) -> None:
+        if self._sqs_ctx is not None:
+            await self._sqs_ctx.__aexit__(None, None, None)
+            self._sqs = None
+            self._sqs_ctx = None
+
+    # ---------- Receive ----------
+
+    async def receive_messages(
+        self,
+        queue_url: Optional[str] = None,
+        max_number_of_messages: int = 1,
+        wait_time_seconds: int = 20,
+        visibility_timeout: Optional[int] = None,
+        attribute_names: Optional[List[str]] = None,
+        message_attribute_names: Optional[List[str]] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Receive up to `max_number_of_messages` from the queue.
+
+        SQS caps MaxNumberOfMessages at 10 per API call, so for larger
+        batch sizes this loops until the requested count is reached or
+        the queue returns no more messages.
+        """
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        sqs = await self._get_client()
+        collected: List[Dict[str, Any]] = []
+        remaining = max_number_of_messages
+
+        first_call = True
+        while remaining > 0:
+            fetch_count = min(remaining, 10)
+
+            params: Dict[str, Any] = {
+                "QueueUrl": url,
+                "MaxNumberOfMessages": fetch_count,
+                "WaitTimeSeconds": wait_time_seconds if first_call else 0,
+                "AttributeNames": attribute_names or ["All"],
+                "MessageAttributeNames": message_attribute_names or ["All"],
+            }
+            if visibility_timeout is not None:
+                params["VisibilityTimeout"] = visibility_timeout
+
+            response = await sqs.receive_message(**params)
+            messages = response.get("Messages", [])
+            first_call = False
+
+            if not messages:
+                break
+
+            collected.extend(messages)
+            remaining -= len(messages)
+
+            if len(messages) < fetch_count:
+                break
+
+        return collected
+
+    # ---------- Send ----------
+
+    async def send_message(
+        self,
+        message_body: Any,
+        queue_url: Optional[str] = None,
+        delay_seconds: int = 0,
+        message_attributes: Optional[Dict[str, Dict]] = None,
+        message_group_id: Optional[str] = None,
+        message_deduplication_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        sqs = await self._get_client()
+        body = json.dumps(message_body) if not isinstance(message_body, str) else message_body
+
+        params: Dict[str, Any] = {
+            "QueueUrl": url,
+            "MessageBody": body,
+            "DelaySeconds": delay_seconds,
+        }
+        if message_attributes:
+            params["MessageAttributes"] = message_attributes
+        if message_group_id:
+            params["MessageGroupId"] = message_group_id
+        if message_deduplication_id:
+            params["MessageDeduplicationId"] = message_deduplication_id
+
+        return await sqs.send_message(**params)
+
+    async def send_message_batch(
+        self,
+        entries: List[Dict[str, Any]],
+        queue_url: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Send messages in batches of up to 10 (SQS limit).
+        Each entry must have at minimum 'Id' and 'MessageBody'.
+        """
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        sqs = await self._get_client()
+        successful = []
+        failed = []
+
+        for i in range(0, len(entries), 10):
+            chunk = entries[i : i + 10]
+            response = await sqs.send_message_batch(QueueUrl=url, Entries=chunk)
+            successful.extend(response.get("Successful", []))
+            failed.extend(response.get("Failed", []))
+
+        return {"Successful": successful, "Failed": failed}
+
+    # ---------- Delete ----------
+
+    async def delete_message(self, receipt_handle: str, queue_url: Optional[str] = None) -> None:
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        sqs = await self._get_client()
+        await sqs.delete_message(QueueUrl=url, ReceiptHandle=receipt_handle)
+
+    async def delete_message_batch(
+        self,
+        entries: List[Dict[str, str]],
+        queue_url: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Delete messages in batches of up to 10.
+        Each entry must have 'Id' and 'ReceiptHandle'.
+        """
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        sqs = await self._get_client()
+        successful = []
+        failed = []
+
+        for i in range(0, len(entries), 10):
+            chunk = entries[i : i + 10]
+            response = await sqs.delete_message_batch(QueueUrl=url, Entries=chunk)
+            successful.extend(response.get("Successful", []))
+            failed.extend(response.get("Failed", []))
+
+        return {"Successful": successful, "Failed": failed}
+
+    # ---------- Queue Management ----------
+
+    async def list_queues(self, prefix: Optional[str] = None) -> List[str]:
+        sqs = await self._get_client()
+        params: Dict[str, Any] = {}
+        if prefix:
+            params["QueueNamePrefix"] = prefix
+
+        response = await sqs.list_queues(**params)
+        return response.get("QueueUrls", [])
+
+    async def get_queue_attributes(
+        self,
+        attribute_names: Optional[List[str]] = None,
+        queue_url: Optional[str] = None,
+    ) -> Dict[str, str]:
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        sqs = await self._get_client()
+        response = await sqs.get_queue_attributes(
+            QueueUrl=url,
+            AttributeNames=attribute_names or ["All"],
+        )
+        return response.get("Attributes", {})
+
+    async def create_queue(
+        self,
+        queue_name: str,
+        attributes: Optional[Dict[str, str]] = None,
+    ) -> str:
+        sqs = await self._get_client()
+        params: Dict[str, Any] = {"QueueName": queue_name}
+        if attributes:
+            params["Attributes"] = attributes
+
+        response = await sqs.create_queue(**params)
+        return response["QueueUrl"]
+
+    async def delete_queue(self, queue_url: Optional[str] = None) -> None:
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        sqs = await self._get_client()
+        await sqs.delete_queue(QueueUrl=url)
+
+    async def purge_queue(self, queue_url: Optional[str] = None) -> None:
+        url = queue_url or self.client.queue_url
+        if not url:
+            raise ValueError("queue_url must be provided either in SQSClient or per-call")
+
+        sqs = await self._get_client()
+        await sqs.purge_queue(QueueUrl=url)

ergon/connector/sqs/connector.py

@@ -0,0 +1,120 @@
+import json
+import logging
+from typing import Any, Dict, List, Optional
+
+from ..connector import Connector
+from ..transaction import Transaction
+from .models import SQSClient, SQSConsumerConfig, SQSProducerConfig
+from .service import SQSService
+
+logger = logging.getLogger(__name__)
+
+
+class SQSConnector(Connector):
+    service: SQSService
+
+    def __init__(
+        self,
+        client: SQSClient,
+        consumer_config: Optional[SQSConsumerConfig] = None,
+        producer_config: Optional[SQSProducerConfig] = None,
+    ) -> None:
+        self.service = SQSService(client)
+        self._consumer_config = consumer_config or SQSConsumerConfig()
+        self._producer_config = producer_config or SQSProducerConfig()
+
+    def fetch_transactions(
+        self,
+        batch_size: int = 1,
+        queue_url: Optional[str] = None,
+        wait_time_seconds: Optional[int] = None,
+        visibility_timeout: Optional[int] = None,
+        attribute_names: Optional[List[str]] = None,
+        message_attribute_names: Optional[List[str]] = None,
+        *args,
+        **kwargs,
+    ) -> List[Transaction]:
+        cfg = self._consumer_config
+
+        resolved_queue_url = queue_url or cfg.queue_url
+        resolved_wait_time = wait_time_seconds if wait_time_seconds is not None else cfg.wait_time_seconds
+        resolved_visibility = visibility_timeout if visibility_timeout is not None else cfg.visibility_timeout
+        resolved_attr_names = attribute_names or cfg.attribute_names
+        resolved_msg_attr_names = message_attribute_names or cfg.message_attribute_names
+
+        messages = self.service.receive_messages(
+            queue_url=resolved_queue_url,
+            max_number_of_messages=batch_size,
+            wait_time_seconds=resolved_wait_time,
+            visibility_timeout=resolved_visibility,
+            attribute_names=resolved_attr_names,
+            message_attribute_names=resolved_msg_attr_names,
+        )
+
+        if not messages:
+            return []
+
+        effective_queue_url = resolved_queue_url or self.service.client.queue_url
+
+        return [
+            Transaction(
+                id=msg["MessageId"],
+                payload=msg,
+                metadata={
+                    "receipt_handle": msg["ReceiptHandle"],
+                    "queue_url": effective_queue_url,
+                },
+            )
+            for msg in messages
+        ]
+
+    def dispatch_transactions(
+        self,
+        transactions: List[Transaction],
+        queue_url: Optional[str] = None,
+        delay_seconds: Optional[int] = None,
+        message_group_id: Optional[str] = None,
+        message_deduplication_id: Optional[str] = None,
+        message_attributes: Optional[Dict[str, Dict]] = None,
+        *args,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        cfg = self._producer_config
+
+        resolved_queue_url = queue_url or cfg.queue_url
+        resolved_delay = delay_seconds if delay_seconds is not None else cfg.delay_seconds
+        resolved_group_id = message_group_id or cfg.message_group_id
+        resolved_dedup_id = message_deduplication_id or cfg.message_deduplication_id
+        resolved_msg_attrs = message_attributes or cfg.message_attributes
+
+        entries: List[Dict[str, Any]] = []
+        for i, txn in enumerate(transactions):
+            body = txn.payload
+            if not isinstance(body, str):
+                body = json.dumps(body)
+
+            entry: Dict[str, Any] = {
+                "Id": str(i),
+                "MessageBody": body,
+                "DelaySeconds": resolved_delay,
+            }
+            if resolved_msg_attrs:
+                entry["MessageAttributes"] = resolved_msg_attrs
+            if resolved_group_id:
+                entry["MessageGroupId"] = resolved_group_id
+
+            txn_dedup_id = txn.metadata.get("message_deduplication_id") if txn.metadata else None
+            dedup_id = txn_dedup_id or resolved_dedup_id
+            if dedup_id:
+                entry["MessageDeduplicationId"] = dedup_id
+
+            entries.append(entry)
+
+        return self.service.send_message_batch(entries=entries, queue_url=resolved_queue_url)
+
+    def get_transactions_count(self, queue_url: Optional[str] = None, *args, **kwargs) -> int:
+        attrs = self.service.get_queue_attributes(
+            attribute_names=["ApproximateNumberOfMessages"],
+            queue_url=queue_url,
+        )
+        return int(attrs.get("ApproximateNumberOfMessages", 0))