jararaca 0.2.37a12__py3-none-any.whl → 0.3.1__py3-none-any.whl

jararaca/messagebus/decorators.py

@@ -2,7 +2,8 @@ import inspect
 from dataclasses import dataclass
 from typing import Any, Awaitable, Callable, Generic, TypeVar, cast
 
-from jararaca.messagebus.types import INHERITS_MESSAGE_CO, Message, MessageOf
+from jararaca.messagebus.message import INHERITS_MESSAGE_CO, Message, MessageOf
+from jararaca.scheduler.decorators import ScheduledAction
 
 DECORATED_FUNC = TypeVar("DECORATED_FUNC", bound=Callable[..., Any])
 DECORATED_CLASS = TypeVar("DECORATED_CLASS", bound=Any)
@@ -45,7 +46,7 @@ class MessageHandler(Generic[INHERITS_MESSAGE_CO]):
 
     @staticmethod
     def get_message_incoming(
-        func: Callable[[MessageOf[Any]], Awaitable[Any]]
+        func: Callable[[MessageOf[Any]], Awaitable[Any]],
     ) -> "MessageHandler[Message] | None":
         if not hasattr(func, MessageHandler.MESSAGE_INCOMING_ATTR):
             return None
@@ -62,7 +63,21 @@ class MessageHandlerData:
     callable: Callable[[MessageOf[Any]], Awaitable[None]]
 
 
+@dataclass(frozen=True)
+class ScheduleDispatchData:
+    timestamp: float
+
+
+@dataclass(frozen=True)
+class ScheduledActionData:
+    spec: ScheduledAction
+    callable: Callable[
+        ..., Awaitable[None]
+    ]  # Callable[[ScheduleDispatchData], Awaitable[None]]
+
+
 MESSAGE_HANDLER_DATA_SET = set[MessageHandlerData]
+SCHEDULED_ACTION_DATA_SET = set[ScheduledActionData]
 
 
 class MessageBusController:
@@ -70,11 +85,16 @@ class MessageBusController:
     MESSAGEBUS_ATTR = "__messagebus__"
 
     def __init__(self) -> None:
-        self.messagebus_factory: Callable[[Any], MESSAGE_HANDLER_DATA_SET] | None = None
+        self.messagebus_factory: (
+            Callable[[Any], tuple[MESSAGE_HANDLER_DATA_SET, SCHEDULED_ACTION_DATA_SET]]
+            | None
+        ) = None
 
     def get_messagebus_factory(
         self,
-    ) -> Callable[[DECORATED_CLASS], MESSAGE_HANDLER_DATA_SET]:
+    ) -> Callable[
+        [DECORATED_CLASS], tuple[MESSAGE_HANDLER_DATA_SET, SCHEDULED_ACTION_DATA_SET]
+    ]:
         if self.messagebus_factory is None:
             raise Exception("MessageBus factory is not set")
         return self.messagebus_factory
@@ -83,33 +103,49 @@ class MessageBusController:
 
         def messagebus_factory(
             instance: DECORATED_CLASS,
-        ) -> MESSAGE_HANDLER_DATA_SET:
+        ) -> tuple[MESSAGE_HANDLER_DATA_SET, SCHEDULED_ACTION_DATA_SET]:
             handlers: MESSAGE_HANDLER_DATA_SET = set()
-            inspect.signature(func)
+
+            schedulers: SCHEDULED_ACTION_DATA_SET = set()
 
             members = inspect.getmembers(func, predicate=inspect.isfunction)
 
             for name, member in members:
-                message_incoming = MessageHandler.get_message_incoming(member)
+                message_handler_decoration = MessageHandler.get_message_incoming(member)
+                scheduled_action_decoration = ScheduledAction.get_scheduled_action(
+                    member
+                )
 
-                if message_incoming is None:
-                    continue
+                if message_handler_decoration is not None:
 
-                if not inspect.iscoroutinefunction(member):
-                    raise Exception(
-                        "Message incoming handler '%s' from '%s.%s' must be a coroutine function"
-                        % (name, func.__module__, func.__qualname__)
-                    )
+                    if not inspect.iscoroutinefunction(member):
+                        raise Exception(
+                            "Message incoming handler '%s' from '%s.%s' must be a coroutine function"
+                            % (name, func.__module__, func.__qualname__)
+                        )
 
-                handlers.add(
-                    MessageHandlerData(
-                        message_type=message_incoming.message_type,
-                        spec=message_incoming,
-                        callable=getattr(instance, name),
+                    handlers.add(
+                        MessageHandlerData(
+                            message_type=message_handler_decoration.message_type,
+                            spec=message_handler_decoration,
+                            callable=getattr(instance, name),
+                        )
+                    )
+                elif scheduled_action_decoration is not None:
+                    if not inspect.iscoroutinefunction(member):
+                        raise Exception(
+                            "Scheduled action handler '%s' from '%s.%s' must be a coroutine function"
+                            % (name, func.__module__, func.__qualname__)
+                        )
+
+                    schedulers.add(
+                        ScheduledActionData(
+                            spec=scheduled_action_decoration,
+                            callable=getattr(instance, name),
+                        )
                     )
-                )
 
-            return handlers
+            return handlers, schedulers
 
         self.messagebus_factory = messagebus_factory
 

jararaca/messagebus/interceptors/aiopika_publisher_interceptor.py

@@ -1,57 +1,73 @@
+import logging
 from contextlib import asynccontextmanager
-from typing import Any, AsyncContextManager, AsyncGenerator, Protocol
+from datetime import datetime, timedelta
+from datetime import tzinfo as _TzInfo
+from typing import Any, AsyncGenerator
 
 import aio_pika
 from aio_pika.abc import AbstractConnection
 from pydantic import BaseModel
 
-from jararaca.messagebus.publisher import MessagePublisher, provide_message_publisher
-from jararaca.microservice import AppContext, AppInterceptor
+from jararaca.broker_backend import MessageBrokerBackend
+from jararaca.messagebus.interceptors.publisher_interceptor import (
+    MessageBusConnectionFactory,
+)
+from jararaca.messagebus.publisher import IMessage, MessagePublisher
+from jararaca.scheduler.types import DelayedMessageData
 
 
-class MessageBusConnectionFactory(Protocol):
-
-    def provide_connection(self) -> AsyncContextManager[MessagePublisher]: ...
-
-
-class MessageBusPublisherInterceptor(AppInterceptor):
+class AIOPikaMessagePublisher(MessagePublisher):
 
     def __init__(
         self,
-        connection_factory: MessageBusConnectionFactory,
-        connection_name: str = "default",
+        channel: aio_pika.abc.AbstractChannel,
+        exchange_name: str,
+        message_broker_backend: MessageBrokerBackend | None = None,
     ):
-        self.connection_factory = connection_factory
-        self.connection_name = connection_name
-
-    @asynccontextmanager
-    async def intercept(self, app_context: AppContext) -> AsyncGenerator[None, None]:
-        if app_context.context_type == "websocket":
-            yield
-            return
-
-        async with self.connection_factory.provide_connection() as connection:
-            with provide_message_publisher(self.connection_name, connection):
-                yield
-
-
-class AIOPikaMessagePublisher(MessagePublisher):
 
-    def __init__(self, channel: aio_pika.abc.AbstractChannel, exchange_name: str):
         self.channel = channel
         self.exchange_name = exchange_name
+        self.message_broker_backend = message_broker_backend
 
-    async def publish(self, message: BaseModel, topic: str) -> None:
-        exchange = await self.channel.declare_exchange(
-            self.exchange_name,
-            type=aio_pika.ExchangeType.TOPIC,
-        )
-        routing_key = f"{self.exchange_name}.{topic}."
+    async def publish(self, message: IMessage, topic: str) -> None:
+        exchange = await self.channel.get_exchange(self.exchange_name, ensure=False)
+        if not exchange:
+            logging.warning(f"Exchange {self.exchange_name} not found")
+            return
+        routing_key = f"{topic}."
         await exchange.publish(
             aio_pika.Message(body=message.model_dump_json().encode()),
             routing_key=routing_key,
         )
 
+    async def delay(self, message: IMessage, seconds: int) -> None:
+        if not self.message_broker_backend:
+            raise NotImplementedError(
+                "Delay is not implemented for AIOPikaMessagePublisher"
+            )
+        await self.message_broker_backend.enqueue_delayed_message(
+            DelayedMessageData(
+                message_topic=message.MESSAGE_TOPIC,
+                payload=message.model_dump_json().encode(),
+                dispatch_time=int(
+                    (datetime.now(tz=None) + timedelta(seconds=seconds)).timestamp()
+                ),
+            )
+        )
+
+    async def schedule(self, message: IMessage, when: datetime, tz: _TzInfo) -> None:
+        if not self.message_broker_backend:
+            raise NotImplementedError(
+                "Schedule is not implemented for AIOPikaMessagePublisher"
+            )
+        await self.message_broker_backend.enqueue_delayed_message(
+            DelayedMessageData(
+                message_topic=message.MESSAGE_TOPIC,
+                payload=message.model_dump_json().encode(),
+                dispatch_time=int(when.timestamp()),
+            )
+        )
+
 
 class GenericPoolConfig(BaseModel):
     max_size: int
@@ -65,10 +81,11 @@ class AIOPikaConnectionFactory(MessageBusConnectionFactory):
         exchange: str,
         connection_pool_config: GenericPoolConfig | None = None,
         channel_pool_config: GenericPoolConfig | None = None,
+        message_broker_backend: MessageBrokerBackend | None = None,
     ):
         self.url = url
         self.exchange = exchange
-
+        self.message_broker_backend = message_broker_backend
         self.connection_pool: aio_pika.pool.Pool[AbstractConnection] | None = None
         self.channel_pool: aio_pika.pool.Pool[aio_pika.abc.AbstractChannel] | None = (
             None
@@ -124,7 +141,11 @@ class AIOPikaConnectionFactory(MessageBusConnectionFactory):
             await tx.select()
 
             try:
-                yield AIOPikaMessagePublisher(channel, exchange_name=self.exchange)
+                yield AIOPikaMessagePublisher(
+                    channel,
+                    exchange_name=self.exchange,
+                    message_broker_backend=self.message_broker_backend,
+                )
                 await tx.commit()
             except Exception as e:
                 await tx.rollback()

jararaca/messagebus/interceptors/publisher_interceptor.py

@@ -0,0 +1,34 @@
+from contextlib import asynccontextmanager
+from typing import AsyncContextManager, AsyncGenerator, Protocol
+
+from jararaca.broker_backend import MessageBrokerBackend
+from jararaca.messagebus.publisher import MessagePublisher, provide_message_publisher
+from jararaca.microservice import AppContext, AppInterceptor
+
+
+class MessageBusConnectionFactory(Protocol):
+
+    def provide_connection(self) -> AsyncContextManager[MessagePublisher]: ...
+
+
+class MessageBusPublisherInterceptor(AppInterceptor):
+
+    def __init__(
+        self,
+        connection_factory: MessageBusConnectionFactory,
+        connection_name: str = "default",
+        message_scheduler: MessageBrokerBackend | None = None,
+    ):
+        self.connection_factory = connection_factory
+        self.connection_name = connection_name
+        self.message_scheduler = message_scheduler
+
+    @asynccontextmanager
+    async def intercept(self, app_context: AppContext) -> AsyncGenerator[None, None]:
+        if app_context.context_type == "websocket":
+            yield
+            return
+
+        async with self.connection_factory.provide_connection() as connection:
+            with provide_message_publisher(self.connection_name, connection):
+                yield

jararaca/messagebus/message.py

@@ -0,0 +1,27 @@
+from datetime import datetime, tzinfo
+from typing import Generic, Protocol, TypeVar
+
+from jararaca.messagebus.publisher import IMessage, use_publisher
+
+
+class Message(IMessage):
+
+    async def publish(self) -> None:
+        task_publisher = use_publisher()
+        await task_publisher.publish(self, self.MESSAGE_TOPIC)
+
+    async def delay(self, seconds: int) -> None:
+        task_publisher = use_publisher()
+        await task_publisher.delay(self, seconds)
+
+    async def schedule(self, when: datetime, tz: tzinfo) -> None:
+        task_publisher = use_publisher()
+        await task_publisher.schedule(self, when, tz)
+
+
+INHERITS_MESSAGE_CO = TypeVar("INHERITS_MESSAGE_CO", bound=Message, covariant=True)
+
+
+class MessageOf(Protocol, Generic[INHERITS_MESSAGE_CO]):
+
+    def payload(self) -> INHERITS_MESSAGE_CO: ...

jararaca/messagebus/publisher.py

@@ -1,12 +1,41 @@
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
-from typing import Any, Generator, Protocol
+from datetime import datetime, tzinfo
+from typing import Any, ClassVar, Generator, Literal, Protocol
 
 from pydantic import BaseModel
 
 
+class IMessage(BaseModel):
+    """
+    Base class for messages representing tasks.
+    A Task is a message that represents a unit of work to be done.
+    It is published to a TaskPublisher and consumed by a TaskHandler, wrapped in TaskData.
+    Note: A Task is not an Event.
+    """
+
+    MESSAGE_TOPIC: ClassVar[str] = "__unset__"
+
+    MESSAGE_TYPE: ClassVar[Literal["task", "event"]] = "task"
+
+
 class MessagePublisher(Protocol):
-    async def publish(self, message: BaseModel, topic: str) -> None:
+    async def publish(self, message: IMessage, topic: str) -> None:
+        raise NotImplementedError()
+
+    async def delay(self, message: IMessage, seconds: int) -> None:
+        """
+        Delay the message for a given number of seconds.
+        """
+
+        raise NotImplementedError()
+
+    async def schedule(
+        self, message: IMessage, when: datetime, timezone: tzinfo
+    ) -> None:
+        """
+        Schedule the message for a given datetime.
+        """
         raise NotImplementedError()
 
 

jararaca/messagebus/worker.py

@@ -24,8 +24,9 @@ from jararaca.messagebus.decorators import (
     MessageHandler,
     MessageHandlerData,
 )
-from jararaca.messagebus.types import Message, MessageOf
+from jararaca.messagebus.message import Message, MessageOf
 from jararaca.microservice import MessageBusAppContext, Microservice
+from jararaca.utils.rabbitmq_utils import RabbitmqUtils
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +34,6 @@ logger = logging.getLogger(__name__)
 @dataclass
 class AioPikaWorkerConfig:
     url: str
-    queue: str
     exchange: str
     prefetch_count: int
 
@@ -99,32 +99,28 @@ class AioPikaMicroserviceConsumer:
 
         await channel.set_qos(prefetch_count=self.config.prefetch_count)
 
-        main_x = await channel.declare_exchange(self.config.exchange, type="topic")
-
-        dlx = await channel.declare_exchange("dlx", type="direct")
-
-        dlq = await channel.declare_queue("dlq")
+        main_ex = await RabbitmqUtils.declare_main_exchange(
+            channel=channel, exchange_name=self.config.exchange
+        )
 
-        await dlq.bind(dlx, routing_key="dlq")
+        dlx, dlq = await RabbitmqUtils.delcare_dl_kit(channel=channel)
 
         for handler in self.message_handler_set:
 
-            queue_name = f"{self.config.exchange}.{handler.message_type.MESSAGE_TOPIC}.{handler.callable.__module__}.{handler.callable.__qualname__}"
-            routing_key = (
-                f"{self.config.exchange}.{handler.message_type.MESSAGE_TOPIC}.#"
-            )
+            queue_name = f"{handler.message_type.MESSAGE_TOPIC}.{handler.callable.__module__}.{handler.callable.__qualname__}"
+            routing_key = f"{handler.message_type.MESSAGE_TOPIC}.#"
 
             self.incoming_map[queue_name] = handler
 
             queue = await channel.declare_queue(
                 queue_name,
                 arguments={
-                    "x-dead-letter-exchange": "dlx",
-                    "x-dead-letter-routing-key": "dlq",
+                    "x-dead-letter-exchange": dlx.name,
+                    "x-dead-letter-routing-key": dlq.name,
                 },
             )
 
-            await queue.bind(exchange=self.config.exchange, routing_key=routing_key)
+            await queue.bind(exchange=main_ex, routing_key=routing_key)
 
             await queue.consume(
                 callback=MessageHandlerCallback(
@@ -343,7 +339,7 @@ class MessageBusWorker:
                 instance: Any = self.container.get_by_type(instance_type)
 
                 factory = controller.get_messagebus_factory()
-                handlers = factory(instance)
+                handlers, _ = factory(instance)
 
                 message_handler_data_map: dict[str, MessageHandlerData] = {}