jararaca 0.3.9__py3-none-any.whl → 0.3.11__py3-none-any.whl

jararaca/core/uow.py

@@ -2,13 +2,14 @@ from contextlib import asynccontextmanager
 from typing import AsyncGenerator, Sequence
 
 from jararaca.microservice import (
-    AppContext,
     AppInterceptor,
+    AppTransactionContext,
     Container,
     Microservice,
     provide_app_context,
     provide_container,
 )
+from jararaca.reflect.metadata import provide_metadata
 
 
 class ContainerInterceptor(AppInterceptor):
@@ -17,7 +18,9 @@ class ContainerInterceptor(AppInterceptor):
         self.container = container
 
     @asynccontextmanager
-    async def intercept(self, app_context: AppContext) -> AsyncGenerator[None, None]:
+    async def intercept(
+        self, app_context: AppTransactionContext
+    ) -> AsyncGenerator[None, None]:
 
         with provide_app_context(app_context), provide_container(self.container):
             yield None
@@ -49,18 +52,20 @@ class UnitOfWorkContextProvider:
         return interceptors
 
     @asynccontextmanager
-    async def __call__(self, app_context: AppContext) -> AsyncGenerator[None, None]:
+    async def __call__(
+        self, app_context: AppTransactionContext
+    ) -> AsyncGenerator[None, None]:
 
         app_interceptors = self.factory_app_interceptors()
+        with provide_metadata(app_context.controller_member_reflect.metadata):
+            ctxs = [self.container_interceptor.intercept(app_context)] + [
+                interceptor.intercept(app_context) for interceptor in app_interceptors
+            ]
 
-        ctxs = [self.container_interceptor.intercept(app_context)] + [
-            interceptor.intercept(app_context) for interceptor in app_interceptors
-        ]
+            for ctx in ctxs:
+                await ctx.__aenter__()
 
-        for ctx in ctxs:
-            await ctx.__aenter__()
-
-        yield None
+            yield None
 
-        for ctx in reversed(ctxs):
-            await ctx.__aexit__(None, None, None)
+            for ctx in reversed(ctxs):
+                await ctx.__aexit__(None, None, None)

jararaca/messagebus/decorators.py

@@ -3,10 +3,14 @@ from dataclasses import dataclass
 from typing import Any, Awaitable, Callable, Generic, TypeVar, cast
 
 from jararaca.messagebus.message import INHERITS_MESSAGE_CO, Message, MessageOf
-from jararaca.scheduler.decorators import ScheduledAction
+from jararaca.reflect.controller_inspect import (
+    ControllerMemberReflect,
+    inspect_controller,
+)
+from jararaca.scheduler.decorators import ScheduledAction, ScheduledActionData
 
 DECORATED_FUNC = TypeVar("DECORATED_FUNC", bound=Callable[..., Any])
-DECORATED_CLASS = TypeVar("DECORATED_CLASS", bound=Any)
+DECORATED_T = TypeVar("DECORATED_T", bound=Any)
 
 
 class MessageHandler(Generic[INHERITS_MESSAGE_CO]):
@@ -20,6 +24,7 @@ class MessageHandler(Generic[INHERITS_MESSAGE_CO]):
         exception_handler: Callable[[BaseException], None] | None = None,
         nack_on_exception: bool = False,
         auto_ack: bool = True,
+        name: str | None = None,
     ) -> None:
         self.message_type = message
 
@@ -27,6 +32,7 @@ class MessageHandler(Generic[INHERITS_MESSAGE_CO]):
         self.exception_handler = exception_handler
         self.requeue_on_exception = nack_on_exception
         self.auto_ack = auto_ack
+        self.name = name
 
     def __call__(
         self, func: Callable[[Any, MessageOf[INHERITS_MESSAGE_CO]], Awaitable[None]]
@@ -60,7 +66,8 @@ class MessageHandler(Generic[INHERITS_MESSAGE_CO]):
 class MessageHandlerData:
     message_type: type[Any]
     spec: MessageHandler[Message]
-    callable: Callable[[MessageOf[Any]], Awaitable[None]]
+    instance_callable: Callable[[MessageOf[Any]], Awaitable[None]]
+    controller_member: ControllerMemberReflect
 
 
 @dataclass(frozen=True)
@@ -68,16 +75,9 @@ class ScheduleDispatchData:
     timestamp: float
 
 
-@dataclass(frozen=True)
-class ScheduledActionData:
-    spec: ScheduledAction
-    callable: Callable[
-        ..., Awaitable[None]
-    ]  # Callable[[ScheduleDispatchData], Awaitable[None]]
-
+SCHEDULED_ACTION_DATA_SET = set[ScheduledActionData]
 
 MESSAGE_HANDLER_DATA_SET = set[MessageHandlerData]
-SCHEDULED_ACTION_DATA_SET = set[ScheduledActionData]
 
 
 class MessageBusController:
@@ -93,55 +93,60 @@ class MessageBusController:
     def get_messagebus_factory(
         self,
     ) -> Callable[
-        [DECORATED_CLASS], tuple[MESSAGE_HANDLER_DATA_SET, SCHEDULED_ACTION_DATA_SET]
+        [DECORATED_T], 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
 
-    def __call__(self, func: type[DECORATED_CLASS]) -> type[DECORATED_CLASS]:
+    def __call__(self, cls_t: type[DECORATED_T]) -> type[DECORATED_T]:
 
         def messagebus_factory(
-            instance: DECORATED_CLASS,
+            instance: DECORATED_T,
         ) -> tuple[MESSAGE_HANDLER_DATA_SET, SCHEDULED_ACTION_DATA_SET]:
             handlers: MESSAGE_HANDLER_DATA_SET = set()
 
             schedulers: SCHEDULED_ACTION_DATA_SET = set()
 
-            members = inspect.getmembers(func, predicate=inspect.isfunction)
+            _, members = inspect_controller(cls_t)
 
-            for name, member in members:
-                message_handler_decoration = MessageHandler.get_message_incoming(member)
+            for name, member in members.items():
+                message_handler_decoration = MessageHandler.get_message_incoming(
+                    member.member_function
+                )
                 scheduled_action_decoration = ScheduledAction.get_scheduled_action(
-                    member
+                    member.member_function
                 )
 
                 if message_handler_decoration is not None:
 
-                    if not inspect.iscoroutinefunction(member):
+                    if not inspect.iscoroutinefunction(member.member_function):
                         raise Exception(
                             "Message incoming handler '%s' from '%s.%s' must be a coroutine function"
-                            % (name, func.__module__, func.__qualname__)
+                            % (name, cls_t.__module__, cls_t.__qualname__)
                         )
 
                     handlers.add(
                         MessageHandlerData(
                             message_type=message_handler_decoration.message_type,
                             spec=message_handler_decoration,
-                            callable=getattr(instance, name),
+                            instance_callable=getattr(instance, name),
+                            controller_member=member,
                         )
                     )
                 elif scheduled_action_decoration is not None:
-                    if not inspect.iscoroutinefunction(member):
+                    if not inspect.iscoroutinefunction(member.member_function):
                         raise Exception(
                             "Scheduled action handler '%s' from '%s.%s' must be a coroutine function"
-                            % (name, func.__module__, func.__qualname__)
+                            % (name, cls_t.__module__, cls_t.__qualname__)
                         )
+                    instance_callable = getattr(instance, name)
 
                     schedulers.add(
                         ScheduledActionData(
+                            controller_member=member,
                             spec=scheduled_action_decoration,
-                            callable=getattr(instance, name),
+                            callable=instance_callable,
                         )
                     )
 
@@ -149,19 +154,17 @@ class MessageBusController:
 
         self.messagebus_factory = messagebus_factory
 
-        MessageBusController.register(func, self)
+        MessageBusController.register(cls_t, self)
 
-        return func
+        return cls_t
 
     @staticmethod
-    def register(
-        func: type[DECORATED_CLASS], messagebus: "MessageBusController"
-    ) -> None:
+    def register(func: type[DECORATED_T], messagebus: "MessageBusController") -> None:
 
         setattr(func, MessageBusController.MESSAGEBUS_ATTR, messagebus)
 
     @staticmethod
-    def get_messagebus(func: type[DECORATED_CLASS]) -> "MessageBusController | None":
+    def get_messagebus(func: type[DECORATED_T]) -> "MessageBusController | None":
         if not hasattr(func, MessageBusController.MESSAGEBUS_ATTR):
             return None
 

jararaca/messagebus/interceptors/aiopika_publisher_interceptor.py

@@ -15,6 +15,8 @@ from jararaca.messagebus.interceptors.publisher_interceptor import (
 from jararaca.messagebus.publisher import IMessage, MessagePublisher
 from jararaca.scheduler.types import DelayedMessageData
 
+logger = logging.getLogger(__name__)
+
 
 class AIOPikaMessagePublisher(MessagePublisher):
 
@@ -28,8 +30,13 @@ class AIOPikaMessagePublisher(MessagePublisher):
         self.channel = channel
         self.exchange_name = exchange_name
         self.message_broker_backend = message_broker_backend
+        self.staged_delayed_messages: list[DelayedMessageData] = []
+        self.staged_messages: list[IMessage] = []
 
     async def publish(self, message: IMessage, topic: str) -> None:
+        self.staged_messages.append(message)
+
+    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")
@@ -45,7 +52,7 @@ class AIOPikaMessagePublisher(MessagePublisher):
             raise NotImplementedError(
                 "Delay is not implemented for AIOPikaMessagePublisher"
             )
-        await self.message_broker_backend.enqueue_delayed_message(
+        self.staged_delayed_messages.append(
             DelayedMessageData(
                 message_topic=message.MESSAGE_TOPIC,
                 payload=message.model_dump_json().encode(),
@@ -60,7 +67,7 @@ class AIOPikaMessagePublisher(MessagePublisher):
             raise NotImplementedError(
                 "Schedule is not implemented for AIOPikaMessagePublisher"
             )
-        await self.message_broker_backend.enqueue_delayed_message(
+        self.staged_delayed_messages.append(
             DelayedMessageData(
                 message_topic=message.MESSAGE_TOPIC,
                 payload=message.model_dump_json().encode(),
@@ -68,6 +75,27 @@ class AIOPikaMessagePublisher(MessagePublisher):
             )
         )
 
+    async def flush(self) -> None:
+        for message in self.staged_messages:
+            logger.debug(
+                f"Publishing message {message.MESSAGE_TOPIC} with payload: {message.model_dump_json()}"
+            )
+            await self._publish(message, message.MESSAGE_TOPIC)
+
+        if len(self.staged_delayed_messages) > 0:
+            if not self.message_broker_backend:
+                raise NotImplementedError(
+                    "MessageBrokerBackend is required to publish delayed messages"
+                )
+
+            for delayed_message in self.staged_delayed_messages:
+                logger.debug(
+                    f"Scheduling delayed message {delayed_message.message_topic} with payload: {delayed_message.payload.decode()}"
+                )
+                await self.message_broker_backend.enqueue_delayed_message(
+                    delayed_message
+                )
+
 
 class GenericPoolConfig(BaseModel):
     max_size: int

jararaca/messagebus/interceptors/publisher_interceptor.py

@@ -3,7 +3,7 @@ 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
+from jararaca.microservice import AppInterceptor, AppTransactionContext
 
 
 class MessageBusConnectionFactory(Protocol):
@@ -24,11 +24,15 @@ class MessageBusPublisherInterceptor(AppInterceptor):
         self.message_scheduler = message_scheduler
 
     @asynccontextmanager
-    async def intercept(self, app_context: AppContext) -> AsyncGenerator[None, None]:
-        if app_context.context_type == "websocket":
+    async def intercept(
+        self, app_context: AppTransactionContext
+    ) -> AsyncGenerator[None, None]:
+        if app_context.transaction_data.context_type == "websocket":
             yield
             return
 
         async with self.connection_factory.provide_connection() as connection:
             with provide_message_publisher(self.connection_name, connection):
                 yield
+
+                await connection.flush()

jararaca/messagebus/publisher.py

@@ -1,7 +1,8 @@
+from abc import ABC, abstractmethod
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
 from datetime import datetime, tzinfo
-from typing import Any, ClassVar, Generator, Literal, Protocol
+from typing import Any, ClassVar, Generator, Literal
 
 from pydantic import BaseModel
 
@@ -19,24 +20,31 @@ class IMessage(BaseModel):
     MESSAGE_TYPE: ClassVar[Literal["task", "event"]] = "task"
 
 
-class MessagePublisher(Protocol):
+class MessagePublisher(ABC):
+    @abstractmethod
     async def publish(self, message: IMessage, topic: str) -> None:
-        raise NotImplementedError()
+        pass
 
+    @abstractmethod
     async def delay(self, message: IMessage, seconds: int) -> None:
         """
         Delay the message for a given number of seconds.
         """
 
-        raise NotImplementedError()
-
+    @abstractmethod
     async def schedule(
         self, message: IMessage, when: datetime, timezone: tzinfo
     ) -> None:
         """
         Schedule the message for a given datetime.
         """
-        raise NotImplementedError()
+
+    @abstractmethod
+    async def flush(self) -> None:
+        """
+        Publish all messages that have been delayed or scheduled.
+        This is typically called at the end of a request or task processing.
+        """
 
 
 message_publishers_ctx = ContextVar[dict[str, MessagePublisher]](