jararaca 0.3.11a16__py3-none-any.whl → 0.4.0a19__py3-none-any.whl

jararaca/messagebus/interceptors/aiopika_publisher_interceptor.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 import logging
 from contextlib import asynccontextmanager
 from datetime import datetime, timedelta
@@ -5,10 +9,12 @@ from datetime import tzinfo as _TzInfo
 from typing import Any, AsyncGenerator
 
 import aio_pika
+import tenacity
 from aio_pika.abc import AbstractConnection
 from pydantic import BaseModel
 
 from jararaca.broker_backend import MessageBrokerBackend
+from jararaca.messagebus import implicit_headers
 from jararaca.messagebus.interceptors.publisher_interceptor import (
     MessageBusConnectionFactory,
 )
@@ -41,9 +47,13 @@ class AIOPikaMessagePublisher(MessagePublisher):
         if not exchange:
             logging.warning(f"Exchange {self.exchange_name} not found")
             return
-        routing_key = f"{topic}."
+        routing_key = f"{topic}.#"
+
+        implicit_headers_data = implicit_headers.use_implicit_headers()
         await exchange.publish(
-            aio_pika.Message(body=message.model_dump_json().encode()),
+            aio_pika.Message(
+                body=message.model_dump_json().encode(), headers=implicit_headers_data
+            ),
             routing_key=routing_key,
         )
 
@@ -62,7 +72,9 @@ class AIOPikaMessagePublisher(MessagePublisher):
             )
         )
 
-    async def schedule(self, message: IMessage, when: datetime, tz: _TzInfo) -> None:
+    async def schedule(
+        self, message: IMessage, when: datetime, timezone: _TzInfo
+    ) -> None:
         if not self.message_broker_backend:
             raise NotImplementedError(
                 "Schedule is not implemented for AIOPikaMessagePublisher"
@@ -101,6 +113,18 @@ class GenericPoolConfig(BaseModel):
     max_size: int
 
 
+default_retry = tenacity.retry(
+    wait=tenacity.wait_exponential_jitter(initial=1, max=60),
+    before_sleep=tenacity.before_sleep_log(logger, logging.WARNING),
+)
+
+default_retry_channel = tenacity.retry(
+    wait=tenacity.wait_exponential_jitter(initial=1, max=60),
+    before_sleep=tenacity.before_sleep_log(logger, logging.WARNING),
+    stop=tenacity.stop_after_attempt(5),
+)
+
+
 class AIOPikaConnectionFactory(MessageBusConnectionFactory):
 
     def __init__(
@@ -121,8 +145,9 @@ class AIOPikaConnectionFactory(MessageBusConnectionFactory):
 
         if connection_pool_config:
 
+            @default_retry
             async def get_connection() -> AbstractConnection:
-                return await aio_pika.connect_robust(self.url)
+                return await aio_pika.connect(self.url)
 
             self.connection_pool = aio_pika.pool.Pool[AbstractConnection](
                 get_connection,
@@ -131,6 +156,7 @@ class AIOPikaConnectionFactory(MessageBusConnectionFactory):
 
         if channel_pool_config:
 
+            @default_retry_channel
             async def get_channel() -> aio_pika.abc.AbstractChannel:
                 async with self.acquire_connection() as connection:
                     return await connection.channel(publisher_confirms=False)
@@ -139,10 +165,14 @@ class AIOPikaConnectionFactory(MessageBusConnectionFactory):
                 get_channel, max_size=channel_pool_config.max_size
             )
 
+    @default_retry
+    async def _connect(self) -> AbstractConnection:
+        return await aio_pika.connect(self.url)
+
     @asynccontextmanager
     async def acquire_connection(self) -> AsyncGenerator[AbstractConnection, Any]:
         if not self.connection_pool:
-            async with await aio_pika.connect_robust(self.url) as connection:
+            async with await self._connect() as connection:
                 yield connection
         else:
 
@@ -153,12 +183,33 @@ class AIOPikaConnectionFactory(MessageBusConnectionFactory):
     async def acquire_channel(
         self,
     ) -> AsyncGenerator[aio_pika.abc.AbstractChannel, Any]:
-        if not self.channel_pool:
-            async with self.acquire_connection() as connection:
-                yield await connection.channel(publisher_confirms=False)
-        else:
-            async with self.channel_pool.acquire() as channel:
-                yield channel
+
+        async for attempt in tenacity.AsyncRetrying(
+            wait=tenacity.wait_exponential_jitter(initial=1, max=10),
+            before_sleep=tenacity.before_sleep_log(logger, logging.WARNING),
+            # stop=tenacity.stop_after_attempt(9000),
+        ):
+            with attempt:
+
+                if not self.connection_pool or not self.channel_pool:
+                    async with self.acquire_connection() as connection:
+                        yield await connection.channel(publisher_confirms=False)
+                else:
+
+                    async with self.connection_pool.acquire() as connection:
+                        if not connection.connected.is_set():
+                            await connection.connect()
+
+                        async with connection.channel(
+                            publisher_confirms=False
+                        ) as channel:
+                            yield channel
+                        # await connection.close()
+                if attempt.retry_state.attempt_number > 1:
+                    logger.warning(
+                        "Later successful connection attempt #%d",
+                        attempt.retry_state.attempt_number,
+                    )
 
     @asynccontextmanager
     async def provide_connection(self) -> AsyncGenerator[AIOPikaMessagePublisher, Any]:

jararaca/messagebus/interceptors/message_publisher_collector.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta, tzinfo
+
+from jararaca.messagebus.publisher import IMessage, MessagePublisher
+
+
+@dataclass
+class CollectorDelayedMessageData:
+    message: IMessage
+    when: datetime
+    timezone: tzinfo
+
+
+class MessagePublisherCollector(MessagePublisher):
+
+    def __init__(self) -> None:
+        self.staged_delayed_messages: list[CollectorDelayedMessageData] = []
+        self.staged_messages: list[IMessage] = []
+
+    async def publish(self, message: IMessage, topic: str) -> None:
+        self.staged_messages.append(message)
+
+    async def delay(self, message: IMessage, seconds: int) -> None:
+        self.staged_delayed_messages.append(
+            CollectorDelayedMessageData(
+                message=message,
+                when=datetime.now(UTC) + timedelta(seconds=seconds),
+                timezone=UTC,
+            )
+        )
+
+    async def schedule(
+        self, message: IMessage, when: datetime, timezone: tzinfo
+    ) -> None:
+        self.staged_delayed_messages.append(
+            CollectorDelayedMessageData(
+                message=message,
+                when=when,
+                timezone=timezone,
+            )
+        )
+
+    async def fill(self, publisher: MessagePublisher) -> None:
+        for message in self.staged_messages:
+            await publisher.publish(message, message.MESSAGE_TOPIC)
+
+        for delayed_message in self.staged_delayed_messages:
+            await publisher.schedule(
+                delayed_message.message,
+                delayed_message.when,
+                delayed_message.timezone,
+            )
+
+    def has_messages(self) -> bool:
+        return bool(self.staged_messages or self.staged_delayed_messages)
+
+    async def flush(self) -> None:
+        raise NotImplementedError("I'm just a poor little collector! :(")

jararaca/messagebus/interceptors/publisher_interceptor.py

@@ -1,7 +1,14 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 from contextlib import asynccontextmanager
 from typing import AsyncContextManager, AsyncGenerator, Protocol
 
 from jararaca.broker_backend import MessageBrokerBackend
+from jararaca.messagebus.interceptors.message_publisher_collector import (
+    MessagePublisherCollector,
+)
 from jararaca.messagebus.publisher import MessagePublisher, provide_message_publisher
 from jararaca.microservice import AppInterceptor, AppTransactionContext
 
@@ -18,10 +25,15 @@ class MessageBusPublisherInterceptor(AppInterceptor):
         connection_factory: MessageBusConnectionFactory,
         connection_name: str = "default",
         message_scheduler: MessageBrokerBackend | None = None,
+        *,
+        open_connection_at_end_of_transaction: bool = False,
     ):
         self.connection_factory = connection_factory
         self.connection_name = connection_name
         self.message_scheduler = message_scheduler
+        self.open_connection_at_end_of_transaction = (
+            open_connection_at_end_of_transaction
+        )
 
     @asynccontextmanager
     async def intercept(
@@ -31,8 +43,22 @@ class MessageBusPublisherInterceptor(AppInterceptor):
             yield
             return
 
-        async with self.connection_factory.provide_connection() as connection:
-            with provide_message_publisher(self.connection_name, connection):
+        if self.open_connection_at_end_of_transaction:
+
+            collector = MessagePublisherCollector()
+            with provide_message_publisher(self.connection_name, collector):
                 yield
 
-                await connection.flush()
+            if collector.has_messages():
+                async with self.connection_factory.provide_connection() as connection:
+                    await collector.fill(connection)
+                    await connection.flush()
+                    return
+        else:
+
+            yield
+
+            async with self.connection_factory.provide_connection() as connection:
+                with provide_message_publisher(self.connection_name, connection):
+
+                    await connection.flush()

jararaca/messagebus/message.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 from datetime import datetime, tzinfo
 from typing import Generic, Protocol, TypeVar
 

jararaca/messagebus/publisher.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 from abc import ABC, abstractmethod
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
@@ -19,6 +23,8 @@ class IMessage(BaseModel):
 
     MESSAGE_TYPE: ClassVar[Literal["task", "event"]] = "task"
 
+    MESSAGE_CATEGORY: ClassVar[str] = "uncategorized"
+
 
 class MessagePublisher(ABC):
     @abstractmethod