jararaca 0.3.11a10__py3-none-any.whl → 0.3.11a12__py3-none-any.whl

jararaca/messagebus/worker.py

@@ -2,15 +2,21 @@ import asyncio
 import inspect
 import logging
 import signal
+from abc import ABC
 from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass
+from datetime import UTC, datetime
 from typing import Any, AsyncContextManager, AsyncGenerator, Type, get_origin
+from urllib.parse import parse_qs, urlparse
 
 import aio_pika
 import aio_pika.abc
 import uvloop
+from aio_pika.exceptions import AMQPError, ChannelClosed, ChannelNotFoundEntity
 from pydantic import BaseModel
 
+from jararaca.broker_backend import MessageBrokerBackend
+from jararaca.broker_backend.mapper import get_message_broker_backend_from_url
 from jararaca.core.uow import UnitOfWorkContextProvider
 from jararaca.di import Container
 from jararaca.lifecycle import AppLifecycle
@@ -20,16 +26,20 @@ from jararaca.messagebus.bus_message_controller import (
 )
 from jararaca.messagebus.decorators import (
     MESSAGE_HANDLER_DATA_SET,
+    SCHEDULED_ACTION_DATA_SET,
     MessageBusController,
     MessageHandler,
     MessageHandlerData,
+    ScheduleDispatchData,
 )
 from jararaca.messagebus.message import Message, MessageOf
 from jararaca.microservice import (
     AppTransactionContext,
     MessageBusTransactionData,
     Microservice,
+    SchedulerTransactionData,
 )
+from jararaca.scheduler.decorators import ScheduledActionData
 from jararaca.utils.rabbitmq_utils import RabbitmqUtils
 
 logger = logging.getLogger(__name__)
@@ -80,15 +90,28 @@ class MessageProcessingLocker:
         await asyncio.gather(*self.current_processing_messages_set)
 
 
-class AioPikaMicroserviceConsumer:
+class MessageBusConsumer(ABC):
+
+    async def consume(self) -> None:
+        raise NotImplementedError("consume method not implemented")
+
+    def shutdown(self) -> None: ...
+
+
+class AioPikaMicroserviceConsumer(MessageBusConsumer):
     def __init__(
         self,
+        broker_backend: MessageBrokerBackend,
         config: AioPikaWorkerConfig,
         message_handler_set: MESSAGE_HANDLER_DATA_SET,
+        scheduled_actions: SCHEDULED_ACTION_DATA_SET,
         uow_context_provider: UnitOfWorkContextProvider,
     ):
+
+        self.broker_backend = broker_backend
         self.config = config
         self.message_handler_set = message_handler_set
+        self.scheduled_actions = scheduled_actions
         self.incoming_map: dict[str, MessageHandlerData] = {}
         self.uow_context_provider = uow_context_provider
         self.shutdown_event = asyncio.Event()
@@ -96,40 +119,44 @@ class AioPikaMicroserviceConsumer:
         self.tasks: set[asyncio.Task[Any]] = set()
 
     async def consume(self) -> None:
+
         connection = await aio_pika.connect(self.config.url)
+
         channel = await connection.channel()
+
         await channel.set_qos(prefetch_count=self.config.prefetch_count)
 
-        # Get existing exchange and DL kit
+        # Get existing exchange and queues
         try:
-            main_ex = await RabbitmqUtils.get_main_exchange(
+            exchange = await RabbitmqUtils.get_main_exchange(
                 channel=channel,
                 exchange_name=self.config.exchange,
             )
-            dlx, dlq = await RabbitmqUtils.get_dl_kit(channel=channel)
-        except Exception as e:
-            logger.error(
-                f"Required exchange or queue infrastructure not found. "
+
+            dlx = await RabbitmqUtils.get_dl_exchange(channel=channel)
+            dlq = await RabbitmqUtils.get_dl_queue(channel=channel)
+        except (ChannelNotFoundEntity, ChannelClosed, AMQPError) as e:
+            logger.critical(
+                f"Required exchange or queue infrastructure not found and passive mode is enabled. "
                 f"Please use the declare command first to create the required infrastructure. Error: {e}"
             )
             self.shutdown_event.set()
             return
 
         for handler in self.message_handler_set:
+
             queue_name = f"{handler.message_type.MESSAGE_TOPIC}.{handler.instance_callable.__module__}.{handler.instance_callable.__qualname__}"
             routing_key = f"{handler.message_type.MESSAGE_TOPIC}.#"
 
             self.incoming_map[queue_name] = handler
 
-            # Get existing queue
             try:
-                queue = await RabbitmqUtils.get_worker_v1_queue(
-                    channel=channel,
-                    queue_name=queue_name,
+                queue = await RabbitmqUtils.get_queue(
+                    channel=channel, queue_name=queue_name
                 )
-            except Exception as e:
+            except (ChannelNotFoundEntity, ChannelClosed, AMQPError) as e:
                 logger.error(
-                    f"Worker queue '{queue_name}' not found. "
+                    f"Queue '{queue_name}' not found and passive mode is enabled. "
                     f"Please use the declare command first to create the queue. Error: {e}"
                 )
                 continue
@@ -144,7 +171,36 @@ class AioPikaMicroserviceConsumer:
                 no_ack=handler.spec.auto_ack,
             )
 
-            logger.info(f"Consuming {queue_name}")
+            logger.info(f"Consuming message handler {queue_name}")
+
+        for scheduled_action in self.scheduled_actions:
+
+            queue_name = f"{scheduled_action.callable.__module__}.{scheduled_action.callable.__qualname__}"
+
+            routing_key = queue_name
+
+            try:
+                queue = await RabbitmqUtils.get_queue(
+                    channel=channel, queue_name=queue_name
+                )
+            except (ChannelNotFoundEntity, ChannelClosed, AMQPError) as e:
+                logger.error(
+                    f"Scheduler queue '{queue_name}' not found and passive mode is enabled. "
+                    f"Please use the declare command first to create the queue. Error: {e}"
+                )
+                continue
+
+            await queue.consume(
+                callback=ScheduledMessageHandlerCallback(
+                    consumer=self,
+                    queue_name=queue_name,
+                    routing_key=routing_key,
+                    scheduled_action=scheduled_action,
+                ),
+                no_ack=True,
+            )
+
+            logger.info(f"Consuming scheduler {queue_name}")
 
         await self.shutdown_event.wait()
         logger.info("Worker shutting down")
@@ -158,6 +214,157 @@ class AioPikaMicroserviceConsumer:
         async with self.lock:
             await asyncio.gather(*self.tasks)
 
+    def shutdown(self) -> None:
+        self.shutdown_event.set()
+
+
+def create_message_bus(
+    broker_url: str,
+    broker_backend: MessageBrokerBackend,
+    scheduled_actions: SCHEDULED_ACTION_DATA_SET,
+    message_handler_set: MESSAGE_HANDLER_DATA_SET,
+    uow_context_provider: UnitOfWorkContextProvider,
+) -> MessageBusConsumer:
+
+    parsed_url = urlparse(broker_url)
+
+    if parsed_url.scheme == "amqp" or parsed_url.scheme == "amqps":
+        assert parsed_url.query, "Query string must be set for AMQP URLs"
+
+        query_params: dict[str, list[str]] = parse_qs(parsed_url.query)
+
+        assert "exchange" in query_params, "Exchange must be set in the query string"
+        assert (
+            len(query_params["exchange"]) == 1
+        ), "Exchange must be set in the query string"
+        assert (
+            "prefetch_count" in query_params
+        ), "Prefetch count must be set in the query string"
+        assert (
+            len(query_params["prefetch_count"]) == 1
+        ), "Prefetch count must be set in the query string"
+        assert query_params["prefetch_count"][
+            0
+        ].isdigit(), "Prefetch count must be an integer in the query string"
+        assert query_params["exchange"][0], "Exchange must be set in the query string"
+        assert query_params["prefetch_count"][
+            0
+        ], "Prefetch count must be set in the query string"
+
+        exchange = query_params["exchange"][0]
+        prefetch_count = int(query_params["prefetch_count"][0])
+
+        config = AioPikaWorkerConfig(
+            url=broker_url,
+            exchange=exchange,
+            prefetch_count=prefetch_count,
+        )
+
+        return AioPikaMicroserviceConsumer(
+            config=config,
+            broker_backend=broker_backend,
+            message_handler_set=message_handler_set,
+            scheduled_actions=scheduled_actions,
+            uow_context_provider=uow_context_provider,
+        )
+
+    raise ValueError(
+        f"Unsupported broker URL scheme: {parsed_url.scheme}. Supported schemes are amqp and amqps"
+    )
+
+
+class ScheduledMessageHandlerCallback:
+    def __init__(
+        self,
+        consumer: AioPikaMicroserviceConsumer,
+        queue_name: str,
+        routing_key: str,
+        scheduled_action: ScheduledActionData,
+    ):
+        self.consumer = consumer
+        self.queue_name = queue_name
+        self.routing_key = routing_key
+        self.scheduled_action = scheduled_action
+
+    async def __call__(
+        self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
+    ) -> None:
+
+        if self.consumer.shutdown_event.is_set():
+            return
+
+        async with self.consumer.lock:
+            task = asyncio.create_task(self.handle_message(aio_pika_message))
+            self.consumer.tasks.add(task)
+            task.add_done_callback(self.handle_message_consume_done)
+
+    def handle_message_consume_done(self, task: asyncio.Task[Any]) -> None:
+        self.consumer.tasks.discard(task)
+
+    async def handle_message(
+        self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
+    ) -> None:
+
+        if self.consumer.shutdown_event.is_set():
+            logger.info("Shutdown event set. Rqueuing message")
+            await aio_pika_message.reject(requeue=True)
+
+        sig = inspect.signature(self.scheduled_action.callable)
+        if len(sig.parameters) == 1:
+
+            task = asyncio.create_task(
+                self.run_with_context(
+                    self.scheduled_action,
+                    (ScheduleDispatchData(int(aio_pika_message.body.decode("utf-8"))),),
+                    {},
+                )
+            )
+
+        elif len(sig.parameters) == 0:
+            task = asyncio.create_task(
+                self.run_with_context(
+                    self.scheduled_action,
+                    (),
+                    {},
+                )
+            )
+        else:
+            logger.warning(
+                "Scheduled action '%s' must have exactly one parameter of type ScheduleDispatchData or no parameters"
+                % self.queue_name
+            )
+            return
+
+        self.consumer.tasks.add(task)
+        task.add_done_callback(self.handle_message_consume_done)
+
+        try:
+            await task
+        except Exception as e:
+
+            logger.exception(
+                f"Error processing scheduled action {self.queue_name}: {e}"
+            )
+
+    async def run_with_context(
+        self,
+        scheduled_action: ScheduledActionData,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> None:
+        async with self.consumer.uow_context_provider(
+            AppTransactionContext(
+                controller_member_reflect=scheduled_action.controller_member,
+                transaction_data=SchedulerTransactionData(
+                    scheduled_to=datetime.now(UTC),
+                    cron_expression=scheduled_action.spec.cron,
+                    triggered_at=datetime.now(UTC),
+                ),
+            )
+        ):
+
+            await scheduled_action.callable(*args, **kwargs)
+
 
 class MessageHandlerCallback:
 
@@ -220,13 +427,7 @@ class MessageHandlerCallback:
             await self.handle_reject_message(aio_pika_message)
             return
 
-        handler_data = self.consumer.incoming_map.get(routing_key)
-
-        if handler_data is None:
-            logger.warning("No handler found for topic '%s'" % routing_key)
-            await self.handle_reject_message(aio_pika_message)
-
-            return
+        handler_data = self.message_handler
 
         handler = handler_data.instance_callable
 
@@ -324,9 +525,18 @@ async def none_context() -> AsyncGenerator[None, None]:
 
 
 class MessageBusWorker:
-    def __init__(self, app: Microservice, config: AioPikaWorkerConfig) -> None:
+    def __init__(
+        self,
+        app: Microservice,
+        broker_url: str,
+        backend_url: str,
+        handler_names: set[str] | None = None,
+    ) -> None:
         self.app = app
-        self.config = config
+        self.backend_url = backend_url
+        self.broker_url = broker_url
+        self.handler_names = handler_names
+
         self.container = Container(app)
         self.lifecycle = AppLifecycle(app, self.container)
 
@@ -334,39 +544,46 @@ class MessageBusWorker:
             app=app, container=self.container
         )
 
-        self._consumer: AioPikaMicroserviceConsumer | None = None
+        self._consumer: MessageBusConsumer | None = None
 
     @property
-    def consumer(self) -> AioPikaMicroserviceConsumer:
+    def consumer(self) -> MessageBusConsumer:
         if self._consumer is None:
             raise RuntimeError("Consumer not started")
         return self._consumer
 
-    async def start_async(self, handler_names: set[str] | None = None) -> None:
+    async def start_async(self) -> None:
         all_message_handlers_set: MESSAGE_HANDLER_DATA_SET = set()
+        all_scheduled_actions_set: SCHEDULED_ACTION_DATA_SET = set()
         async with self.lifecycle():
-            for instance_type in self.app.controllers:
-                controller = MessageBusController.get_messagebus(instance_type)
+            for instance_class in self.app.controllers:
+                controller = MessageBusController.get_messagebus(instance_class)
 
                 if controller is None:
                     continue
 
-                instance: Any = self.container.get_by_type(instance_type)
+                instance: Any = self.container.get_by_type(instance_class)
 
                 factory = controller.get_messagebus_factory()
-                handlers, _ = factory(instance)
+                handlers, schedulers = factory(instance)
 
                 message_handler_data_map: dict[str, MessageHandlerData] = {}
-
+                all_scheduled_actions_set.update(schedulers)
                 for handler_data in handlers:
                     message_type = handler_data.spec.message_type
                     topic = message_type.MESSAGE_TOPIC
 
                     # Filter handlers by name if specified
-                    if handler_names is not None and handler_data.spec.name is not None:
-                        if handler_data.spec.name not in handler_names:
+                    if (
+                        self.handler_names is not None
+                        and handler_data.spec.name is not None
+                    ):
+                        if handler_data.spec.name not in self.handler_names:
                             continue
-                    elif handler_names is not None and handler_data.spec.name is None:
+                    elif (
+                        self.handler_names is not None
+                        and handler_data.spec.name is None
+                    ):
                         # Skip handlers without names when filtering is requested
                         continue
 
@@ -382,25 +599,29 @@ class MessageBusWorker:
                     message_handler_data_map[topic] = handler_data
                     all_message_handlers_set.add(handler_data)
 
-            consumer = self._consumer = AioPikaMicroserviceConsumer(
-                config=self.config,
+            broker_backend = get_message_broker_backend_from_url(url=self.backend_url)
+
+            consumer = self._consumer = create_message_bus(
+                broker_url=self.broker_url,
+                broker_backend=broker_backend,
+                scheduled_actions=all_scheduled_actions_set,
                 message_handler_set=all_message_handlers_set,
                 uow_context_provider=self.uow_context_provider,
             )
 
             await consumer.consume()
 
-    def start_sync(self, handler_names: set[str] | None = None) -> None:
+    def start_sync(self) -> None:
 
         def on_shutdown(loop: asyncio.AbstractEventLoop) -> None:
             logger.info("Shutting down")
-            self.consumer.shutdown_event.set()
+            self.consumer.shutdown()
 
         with asyncio.Runner(loop_factory=uvloop.new_event_loop) as runner:
             runner.get_loop().add_signal_handler(
                 signal.SIGINT, on_shutdown, runner.get_loop()
             )
-            runner.run(self.start_async(handler_names=handler_names))
+            runner.run(self.start_async())
 
 
 class AioPikaMessageBusController(BusMessageController):

jararaca/scheduler/{scheduler_v2.py → beat_worker.py}

@@ -4,10 +4,9 @@ import logging
 import signal
 import time
 from abc import ABC, abstractmethod
-from contextlib import asynccontextmanager
 from datetime import UTC, datetime
 from types import FrameType
-from typing import Any, AsyncGenerator
+from typing import Any
 from urllib.parse import parse_qs
 
 import aio_pika
@@ -36,7 +35,7 @@ from jararaca.utils.rabbitmq_utils import RabbitmqUtils
 logger = logging.getLogger(__name__)
 
 
-def extract_scheduled_actions(
+def _extract_scheduled_actions(
     app: Microservice, container: Container, scheduler_names: set[str] | None = None
 ) -> list[ScheduledActionData]:
     scheduled_actions: list[ScheduledActionData] = []
@@ -64,7 +63,7 @@ def extract_scheduled_actions(
 # region Message Broker Dispatcher
 
 
-class MessageBrokerDispatcher(ABC):
+class _MessageBrokerDispatcher(ABC):
 
     @abstractmethod
     async def dispatch_scheduled_action(
@@ -100,7 +99,7 @@ class MessageBrokerDispatcher(ABC):
         pass
 
 
-class RabbitMQBrokerDispatcher(MessageBrokerDispatcher):
+class _RabbitMQBrokerDispatcher(_MessageBrokerDispatcher):
 
     def __init__(self, url: str) -> None:
         self.url = url
@@ -196,7 +195,7 @@ class RabbitMQBrokerDispatcher(MessageBrokerDispatcher):
                 queue_name = ScheduledAction.get_function_id(sched_act_data.callable)
 
                 # Try to get existing queue
-                await RabbitmqUtils.get_scheduler_queue(
+                await RabbitmqUtils.get_scheduled_action_queue(
                     channel=channel,
                     queue_name=queue_name,
                 )
@@ -206,13 +205,13 @@ class RabbitMQBrokerDispatcher(MessageBrokerDispatcher):
         await self.conn_pool.close()
 
 
-def get_message_broker_dispatcher_from_url(url: str) -> MessageBrokerDispatcher:
+def _get_message_broker_dispatcher_from_url(url: str) -> _MessageBrokerDispatcher:
     """
     Factory function to create a message broker instance from a URL.
     Currently, only RabbitMQ is supported.
     """
     if url.startswith("amqp://") or url.startswith("amqps://"):
-        return RabbitMQBrokerDispatcher(url=url)
+        return _RabbitMQBrokerDispatcher(url=url)
     else:
         raise ValueError(f"Unsupported message broker URL: {url}")
 
@@ -220,7 +219,7 @@ def get_message_broker_dispatcher_from_url(url: str) -> MessageBrokerDispatcher:
 # endregion
 
 
-class SchedulerV2:
+class BeatWorker:
 
     def __init__(
         self,
@@ -228,11 +227,11 @@ class SchedulerV2:
         interval: int,
         broker_url: str,
         backend_url: str,
-        scheduler_names: set[str] | None = None,
+        scheduled_action_names: set[str] | None = None,
     ) -> None:
         self.app = app
 
-        self.broker: MessageBrokerDispatcher = get_message_broker_dispatcher_from_url(
+        self.broker: _MessageBrokerDispatcher = _get_message_broker_dispatcher_from_url(
             broker_url
         )
         self.backend: MessageBrokerBackend = get_message_broker_backend_from_url(
@@ -240,7 +239,7 @@ class SchedulerV2:
         )
 
         self.interval = interval
-        self.scheduler_names = scheduler_names
+        self.scheduler_names = scheduled_action_names
         self.container = Container(self.app)
         self.uow_provider = UnitOfWorkContextProvider(app, self.container)
 
@@ -266,7 +265,7 @@ class SchedulerV2:
         """
         async with self.lifecycle():
 
-            scheduled_actions = extract_scheduled_actions(
+            scheduled_actions = _extract_scheduled_actions(
                 self.app, self.container, self.scheduler_names
             )
 
@@ -341,14 +340,3 @@ class SchedulerV2:
 
         await self.backend.dispose()
         await self.broker.dispose()
-
-
-@asynccontextmanager
-async def none_context() -> AsyncGenerator[None, None]:
-    yield
-
-
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)

jararaca/utils/rabbitmq_utils.py

@@ -222,7 +222,7 @@ class RabbitmqUtils:
             raise
 
     @classmethod
-    async def declare_queue(
+    async def declare_worker_queue(
         cls,
         channel: AbstractChannel,
         queue_name: str,
@@ -243,64 +243,13 @@ class RabbitmqUtils:
         )
 
     @classmethod
-    async def get_worker_v1_queue(
+    async def get_scheduled_action_queue(
         cls,
         channel: AbstractChannel,
         queue_name: str,
     ) -> AbstractQueue:
         """
-        Get a worker v1 queue.
-        """
-        try:
-            return await channel.get_queue(queue_name)
-        except ChannelNotFoundEntity as e:
-            logger.error(
-                f"Worker queue '{queue_name}' does not exist. "
-                f"Please use the declare command to create it first. Error: {e}"
-            )
-            raise
-        except ChannelClosed as e:
-            logger.error(
-                f"Channel closed while getting worker queue '{queue_name}'. "
-                f"Error: {e}"
-            )
-            raise
-        except AMQPError as e:
-            logger.error(
-                f"AMQP error while getting worker queue '{queue_name}'. " f"Error: {e}"
-            )
-            raise
-
-    @classmethod
-    async def declare_worker_v1_queue(
-        cls,
-        channel: AbstractChannel,
-        queue_name: str,
-        dlx_name: str,
-        dlq_name: str,
-        passive: bool = False,
-    ) -> AbstractQueue:
-        """
-        Declare a worker v1 queue with custom dead letter exchange and routing key.
-        """
-        return await channel.declare_queue(
-            passive=passive,
-            name=queue_name,
-            arguments={
-                "x-dead-letter-exchange": dlx_name,
-                "x-dead-letter-routing-key": dlq_name,
-            },
-            durable=True,
-        )
-
-    @classmethod
-    async def get_scheduler_queue(
-        cls,
-        channel: AbstractChannel,
-        queue_name: str,
-    ) -> AbstractQueue:
-        """
-        Get a scheduler queue.
+        Get a scheduled action queue.
         """
         try:
             return await channel.get_queue(queue_name)
@@ -324,19 +273,24 @@ class RabbitmqUtils:
             raise
 
     @classmethod
-    async def declare_scheduler_queue(
+    async def declare_scheduled_action_queue(
         cls,
         channel: AbstractChannel,
         queue_name: str,
         passive: bool = False,
     ) -> AbstractQueue:
         """
-        Declare a scheduler queue with simple durable configuration.
+        Declare a scheduled action queue with simple durable configuration.
+        The queue has a max length of 1 to ensure only one scheduled task
+        is processed at a time.
         """
         return await channel.declare_queue(
             name=queue_name,
             durable=True,
             passive=passive,
+            arguments={
+                "x-max-length": 1,
+            },
         )
 
     @classmethod

{jararaca-0.3.11a10.dist-info → jararaca-0.3.11a12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: jararaca
-Version: 0.3.11a10
+Version: 0.3.11a12
 Summary: A simple and fast API framework for Python
 Author: Lucas S
 Author-email: me@luscasleo.dev