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

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, Callable
+from typing import Any
 from urllib.parse import parse_qs
 
 import aio_pika
@@ -25,25 +24,37 @@ from jararaca.core.uow import UnitOfWorkContextProvider
 from jararaca.di import Container
 from jararaca.lifecycle import AppLifecycle
 from jararaca.microservice import Microservice
-from jararaca.scheduler.decorators import ScheduledAction
+from jararaca.scheduler.decorators import (
+    ScheduledAction,
+    ScheduledActionData,
+    get_type_scheduled_actions,
+)
 from jararaca.scheduler.types import DelayedMessageData
+from jararaca.utils.rabbitmq_utils import RabbitmqUtils
 
 logger = logging.getLogger(__name__)
 
-SCHEDULED_ACTION_LIST = list[tuple[Callable[..., Any], "ScheduledAction"]]
-
 
-def extract_scheduled_actions(
-    app: Microservice, container: Container
-) -> SCHEDULED_ACTION_LIST:
-    scheduled_actions: SCHEDULED_ACTION_LIST = []
+def _extract_scheduled_actions(
+    app: Microservice, container: Container, scheduler_names: set[str] | None = None
+) -> list[ScheduledActionData]:
+    scheduled_actions: list[ScheduledActionData] = []
     for controllers in app.controllers:
 
         controller_instance: Any = container.get_by_type(controllers)
 
-        controller_scheduled_actions = ScheduledAction.get_type_scheduled_actions(
-            controller_instance
-        )
+        controller_scheduled_actions = get_type_scheduled_actions(controller_instance)
+
+        # Filter scheduled actions by name if scheduler_names is provided
+        if scheduler_names is not None:
+            filtered_actions = []
+            for action in controller_scheduled_actions:
+                # Include actions that have a name and it's in the provided set
+                if action.spec.name and action.spec.name in scheduler_names:
+                    filtered_actions.append(action)
+                # Skip actions without names when filtering is active
+            controller_scheduled_actions = filtered_actions
+
         scheduled_actions.extend(controller_scheduled_actions)
 
     return scheduled_actions
@@ -52,7 +63,7 @@ def extract_scheduled_actions(
 # region Message Broker Dispatcher
 
 
-class MessageBrokerDispatcher(ABC):
+class _MessageBrokerDispatcher(ABC):
 
     @abstractmethod
     async def dispatch_scheduled_action(
@@ -81,14 +92,14 @@ class MessageBrokerDispatcher(ABC):
         raise NotImplementedError("dispatch_delayed_message() is not implemented yet.")
 
     @abstractmethod
-    async def initialize(self, scheduled_actions: SCHEDULED_ACTION_LIST) -> None:
+    async def initialize(self, scheduled_actions: list[ScheduledActionData]) -> None:
         raise NotImplementedError("initialize() is not implemented yet.")
 
     async def dispose(self) -> None:
         pass
 
 
-class RabbitMQBrokerDispatcher(MessageBrokerDispatcher):
+class _RabbitMQBrokerDispatcher(_MessageBrokerDispatcher):
 
     def __init__(self, url: str) -> None:
         self.url = url
@@ -145,7 +156,7 @@ class RabbitMQBrokerDispatcher(MessageBrokerDispatcher):
 
         logger.info(f"Dispatching message to {action_id} at {timestamp}")
         async with self.channel_pool.acquire() as channel:
-            exchange = await channel.get_exchange(self.exchange)
+            exchange = await RabbitmqUtils.get_main_exchange(channel, self.exchange)
 
             await exchange.publish(
                 aio_pika.Message(body=str(timestamp).encode()),
@@ -163,7 +174,7 @@ class RabbitMQBrokerDispatcher(MessageBrokerDispatcher):
         """
         async with self.channel_pool.acquire() as channel:
 
-            exchange = await channel.get_exchange(self.exchange)
+            exchange = await RabbitmqUtils.get_main_exchange(channel, self.exchange)
             await exchange.publish(
                 aio_pika.Message(
                     body=delayed_message.payload,
@@ -171,32 +182,22 @@ class RabbitMQBrokerDispatcher(MessageBrokerDispatcher):
                 routing_key=f"{delayed_message.message_topic}.",
             )
 
-    async def initialize(self, scheduled_actions: SCHEDULED_ACTION_LIST) -> None:
+    async def initialize(self, scheduled_actions: list[ScheduledActionData]) -> None:
         """
         Initialize the RabbitMQ server.
         This is used to create the exchange and queues for the scheduled actions.
         """
 
         async with self.channel_pool.acquire() as channel:
+            await RabbitmqUtils.get_main_exchange(channel, self.exchange)
 
-            await channel.set_qos(prefetch_count=1)
-
-            await channel.declare_exchange(
-                name=self.exchange,
-                type="topic",
-                durable=True,
-                auto_delete=False,
-            )
+            for sched_act_data in scheduled_actions:
+                queue_name = ScheduledAction.get_function_id(sched_act_data.callable)
 
-            for func, _ in scheduled_actions:
-                queue = await channel.declare_queue(
-                    name=ScheduledAction.get_function_id(func),
-                    durable=True,
-                )
-
-                await queue.bind(
-                    exchange=self.exchange,
-                    routing_key=ScheduledAction.get_function_id(func),
+                # Try to get existing queue
+                await RabbitmqUtils.get_scheduled_action_queue(
+                    channel=channel,
+                    queue_name=queue_name,
                 )
 
     async def dispose(self) -> None:
@@ -204,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}")
 
@@ -218,7 +219,7 @@ def get_message_broker_dispatcher_from_url(url: str) -> MessageBrokerDispatcher:
 # endregion
 
 
-class SchedulerV2:
+class BeatWorker:
 
     def __init__(
         self,
@@ -226,10 +227,11 @@ class SchedulerV2:
         interval: int,
         broker_url: str,
         backend_url: str,
+        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(
@@ -237,6 +239,7 @@ class SchedulerV2:
         )
 
         self.interval = interval
+        self.scheduler_names = scheduled_action_names
         self.container = Container(self.app)
         self.uow_provider = UnitOfWorkContextProvider(app, self.container)
 
@@ -262,19 +265,23 @@ class SchedulerV2:
         """
         async with self.lifecycle():
 
-            scheduled_actions = extract_scheduled_actions(self.app, self.container)
+            scheduled_actions = _extract_scheduled_actions(
+                self.app, self.container, self.scheduler_names
+            )
 
             await self.broker.initialize(scheduled_actions)
 
             await self.run_scheduled_actions(scheduled_actions)
 
     async def run_scheduled_actions(
-        self, scheduled_actions: SCHEDULED_ACTION_LIST
+        self, scheduled_actions: list[ScheduledActionData]
     ) -> None:
 
         while not self.shutdown_event.is_set():
             now = int(time.time())
-            for func, scheduled_action in scheduled_actions:
+            for sched_act_data in scheduled_actions:
+                func = sched_act_data.callable
+                scheduled_action = sched_act_data.spec
                 if self.shutdown_event.is_set():
                     break
 
@@ -333,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/scheduler/decorators.py

@@ -1,5 +1,11 @@
 import inspect
-from typing import Any, Callable, TypeVar, cast
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable, TypeVar, cast
+
+from jararaca.reflect.controller_inspect import (
+    ControllerMemberReflect,
+    inspect_controller,
+)
 
 DECORATED_FUNC = TypeVar("DECORATED_FUNC", bound=Callable[..., Any])
 
@@ -14,6 +20,7 @@ class ScheduledAction:
         exclusive: bool = True,
         timeout: int | None = None,
         exception_handler: Callable[[BaseException], None] | None = None,
+        name: str | None = None,
     ) -> None:
         """
         :param cron: A string representing the cron expression for the scheduled action.
@@ -21,6 +28,7 @@ class ScheduledAction:
         :param exclusive: A boolean indicating if the scheduled action should be executed in one instance of the application. (Requires a distributed lock provided by a backend)
         :param exception_handler: A callable that will be called when an exception is raised during the execution of the scheduled action.
         :param timeout: An integer representing the timeout for the scheduled action in seconds. If the scheduled action takes longer than this time, it will be terminated.
+        :param name: An optional name for the scheduled action, used for filtering which actions to run.
         """
         self.cron = cron
         """
@@ -49,6 +57,11 @@ class ScheduledAction:
         If the scheduled action takes longer than this time, it will be terminated.
         """
 
+        self.name = name
+        """
+        An optional name for the scheduled action, used for filtering which actions to run.
+        """
+
     def __call__(self, func: DECORATED_FUNC) -> DECORATED_FUNC:
         ScheduledAction.register(func, self)
         return func
@@ -66,25 +79,6 @@ class ScheduledAction:
             ScheduledAction, getattr(func, ScheduledAction.SCHEDULED_ACTION_ATTR)
         )
 
-    @staticmethod
-    def get_type_scheduled_actions(
-        instance: Any,
-    ) -> list[tuple[Callable[..., Any], "ScheduledAction"]]:
-
-        members = inspect.getmembers(instance, predicate=inspect.ismethod)
-
-        scheduled_actions: list[tuple[Callable[..., Any], "ScheduledAction"]] = []
-
-        for _, member in members:
-            scheduled_action = ScheduledAction.get_scheduled_action(member)
-
-            if scheduled_action is None:
-                continue
-
-            scheduled_actions.append((member, scheduled_action))
-
-        return scheduled_actions
-
     @staticmethod
     def get_function_id(
         func: Callable[..., Any],
@@ -94,3 +88,44 @@ class ScheduledAction:
         This is used to identify the scheduled action in the message broker.
         """
         return f"{func.__module__}.{func.__qualname__}"
+
+
+@dataclass(frozen=True)
+class ScheduledActionData:
+    spec: ScheduledAction
+    controller_member: ControllerMemberReflect
+    callable: Callable[..., Awaitable[None]]
+
+
+def get_type_scheduled_actions(
+    instance: Any,
+) -> list[ScheduledActionData]:
+
+    _, member_metadata_map = inspect_controller(instance.__class__)
+
+    members = inspect.getmembers(instance, predicate=inspect.ismethod)
+
+    scheduled_actions: list[ScheduledActionData] = []
+
+    for name, member in members:
+        scheduled_action = ScheduledAction.get_scheduled_action(member)
+
+        if scheduled_action is None:
+            continue
+
+        if name not in member_metadata_map:
+            raise Exception(
+                f"Member '{name}' is not a valid controller member in '{instance.__class__.__name__}'"
+            )
+
+        member_metadata = member_metadata_map[name]
+
+        scheduled_actions.append(
+            ScheduledActionData(
+                callable=member,
+                spec=scheduled_action,
+                controller_member=member_metadata,
+            )
+        )
+
+    return scheduled_actions

jararaca/tools/app_config/interceptor.py

@@ -7,9 +7,9 @@ from pydantic import BaseModel
 
 from jararaca.core.providers import Token
 from jararaca.microservice import (
-    AppContext,
     AppInterceptor,
     AppInterceptorWithLifecycle,
+    AppTransactionContext,
     Container,
     Microservice,
 )
@@ -40,7 +40,9 @@ class AppConfigurationInterceptor(AppInterceptor, AppInterceptorWithLifecycle):
         self.config_parser = config_parser
 
     @asynccontextmanager
-    async def intercept(self, app_context: AppContext) -> AsyncGenerator[None, None]:
+    async def intercept(
+        self, app_context: AppTransactionContext
+    ) -> AsyncGenerator[None, None]:
         yield
 
     def instance_basemodels(self, basemodel_type: Type[BaseModel]) -> BaseModel:

jararaca/utils/rabbitmq_utils.py

@@ -1,4 +1,9 @@
+import logging
+
 from aio_pika.abc import AbstractChannel, AbstractExchange, AbstractQueue
+from aio_pika.exceptions import AMQPError, ChannelClosed, ChannelNotFoundEntity
+
+logger = logging.getLogger(__name__)
 
 
 class RabbitmqUtils:
@@ -6,6 +11,39 @@ class RabbitmqUtils:
     DEAD_LETTER_EXCHANGE = "dlx"
     DEAD_LETTER_QUEUE = "dlq"
 
+    # Note: get_worker_v1_queue method is already defined above
+
+    DEAD_LETTER_EXCHANGE = "dlx"
+    DEAD_LETTER_QUEUE = "dlq"
+
+    @classmethod
+    async def get_dl_exchange(cls, channel: AbstractChannel) -> AbstractExchange:
+        """
+        Get the Dead Letter Exchange (DLX) for the given channel.
+        """
+        try:
+            return await channel.get_exchange(
+                cls.DEAD_LETTER_EXCHANGE,
+            )
+        except ChannelNotFoundEntity as e:
+            logger.error(
+                f"Dead Letter Exchange '{cls.DEAD_LETTER_EXCHANGE}' 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 Dead Letter Exchange '{cls.DEAD_LETTER_EXCHANGE}'. "
+                f"Error: {e}"
+            )
+            raise
+        except AMQPError as e:
+            logger.error(
+                f"AMQP error while getting Dead Letter Exchange '{cls.DEAD_LETTER_EXCHANGE}'. "
+                f"Error: {e}"
+            )
+            raise
+
     @classmethod
     async def declare_dl_exchange(
         cls, channel: AbstractChannel, passive: bool
@@ -13,7 +51,7 @@ class RabbitmqUtils:
         """
         Declare a Dead Letter Exchange (DLX) for the given channel.
         """
-        await channel.set_qos(prefetch_count=1)
+
         return await channel.declare_exchange(
             cls.DEAD_LETTER_EXCHANGE,
             passive=passive,
@@ -22,6 +60,34 @@ class RabbitmqUtils:
             auto_delete=False,
         )
 
+    @classmethod
+    async def get_dl_queue(cls, channel: AbstractChannel) -> AbstractQueue:
+        """
+        Get the Dead Letter Queue (DLQ) for the given channel.
+        """
+        try:
+            return await channel.get_queue(
+                cls.DEAD_LETTER_QUEUE,
+            )
+        except ChannelNotFoundEntity as e:
+            logger.error(
+                f"Dead Letter Queue '{cls.DEAD_LETTER_QUEUE}' 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 Dead Letter Queue '{cls.DEAD_LETTER_QUEUE}'. "
+                f"Error: {e}"
+            )
+            raise
+        except AMQPError as e:
+            logger.error(
+                f"AMQP error while getting Dead Letter Queue '{cls.DEAD_LETTER_QUEUE}'. "
+                f"Error: {e}"
+            )
+            raise
+
     @classmethod
     async def declare_dl_queue(
         cls, channel: AbstractChannel, passive: bool
@@ -29,7 +95,7 @@ class RabbitmqUtils:
         """
         Declare a Dead Letter Queue (DLQ) for the given queue.
         """
-        await channel.set_qos(prefetch_count=1)
+
         return await channel.declare_queue(
             cls.DEAD_LETTER_QUEUE,
             durable=True,
@@ -40,6 +106,36 @@ class RabbitmqUtils:
             },
         )
 
+    @classmethod
+    async def get_dl_kit(
+        cls,
+        channel: AbstractChannel,
+    ) -> tuple[AbstractExchange, AbstractQueue]:
+        """
+        Get the Dead Letter Exchange and Queue (DLX and DLQ) for the given channel.
+        """
+        try:
+            dlx = await cls.get_dl_exchange(channel)
+            dlq = await cls.get_dl_queue(channel)
+            return dlx, dlq
+        except ChannelNotFoundEntity as e:
+            logger.error(
+                f"Dead Letter infrastructure does not exist completely. "
+                f"Please use the declare command to create it first. Error: {e}"
+            )
+            raise
+        except ChannelClosed as e:
+            logger.error(
+                f"Channel closed while getting Dead Letter infrastructure. "
+                f"Error: {e}"
+            )
+            raise
+        except AMQPError as e:
+            logger.error(
+                f"AMQP error while getting Dead Letter infrastructure. " f"Error: {e}"
+            )
+            raise
+
     @classmethod
     async def declare_dl_kit(
         cls,
@@ -54,6 +150,33 @@ class RabbitmqUtils:
         await dlq.bind(dlx, routing_key=cls.DEAD_LETTER_EXCHANGE)
         return dlx, dlq
 
+    @classmethod
+    async def get_main_exchange(
+        cls, channel: AbstractChannel, exchange_name: str
+    ) -> AbstractExchange:
+        """
+        Get the main exchange for the given channel.
+        """
+        try:
+            return await channel.get_exchange(exchange_name)
+        except ChannelNotFoundEntity as e:
+            logger.error(
+                f"Exchange '{exchange_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 exchange '{exchange_name}'. "
+                f"Error: {e}"
+            )
+            raise
+        except AMQPError as e:
+            logger.error(
+                f"AMQP error while getting exchange '{exchange_name}'. " f"Error: {e}"
+            )
+            raise
+
     @classmethod
     async def declare_main_exchange(
         cls, channel: AbstractChannel, exchange_name: str, passive: bool
@@ -61,7 +184,7 @@ class RabbitmqUtils:
         """
         Declare a main exchange for the given channel.
         """
-        await channel.set_qos(prefetch_count=1)
+
         return await channel.declare_exchange(
             exchange_name,
             passive=passive,
@@ -71,7 +194,35 @@ class RabbitmqUtils:
         )
 
     @classmethod
-    async def declare_queue(
+    async def get_queue(
+        cls,
+        channel: AbstractChannel,
+        queue_name: str,
+    ) -> AbstractQueue:
+        """
+        Get a queue with the given name.
+        """
+        try:
+            return await channel.get_queue(queue_name)
+        except ChannelNotFoundEntity as e:
+            logger.error(
+                f"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 queue '{queue_name}'. " f"Error: {e}"
+            )
+            raise
+        except AMQPError as e:
+            logger.error(
+                f"AMQP error while getting queue '{queue_name}'. " f"Error: {e}"
+            )
+            raise
+
+    @classmethod
+    async def declare_worker_queue(
         cls,
         channel: AbstractChannel,
         queue_name: str,
@@ -80,7 +231,7 @@ class RabbitmqUtils:
         """
         Declare a queue with the given name and properties.
         """
-        await channel.set_qos(prefetch_count=1)
+
         return await channel.declare_queue(
             queue_name,
             passive=passive,
@@ -90,3 +241,106 @@ class RabbitmqUtils:
                 "x-dead-letter-routing-key": cls.DEAD_LETTER_EXCHANGE,
             },
         )
+
+    @classmethod
+    async def get_scheduled_action_queue(
+        cls,
+        channel: AbstractChannel,
+        queue_name: str,
+    ) -> AbstractQueue:
+        """
+        Get a scheduled action queue.
+        """
+        try:
+            return await channel.get_queue(queue_name)
+        except ChannelNotFoundEntity as e:
+            logger.error(
+                f"Scheduler 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 scheduler queue '{queue_name}'. "
+                f"Error: {e}"
+            )
+            raise
+        except AMQPError as e:
+            logger.error(
+                f"AMQP error while getting scheduler queue '{queue_name}'. "
+                f"Error: {e}"
+            )
+            raise
+
+    @classmethod
+    async def declare_scheduled_action_queue(
+        cls,
+        channel: AbstractChannel,
+        queue_name: str,
+        passive: bool = False,
+    ) -> AbstractQueue:
+        """
+        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
+    async def delete_exchange(
+        cls,
+        channel: AbstractChannel,
+        exchange_name: str,
+        if_unused: bool = False,
+    ) -> None:
+        """
+        Delete an exchange.
+        """
+        try:
+            await channel.exchange_delete(
+                exchange_name=exchange_name,
+                if_unused=if_unused,
+            )
+        except ChannelNotFoundEntity:
+            # Exchange might not exist, which is fine
+            logger.info(
+                f"Exchange '{exchange_name}' does not exist, nothing to delete."
+            )
+        except ChannelClosed as e:
+            logger.warning(
+                f"Channel closed while deleting exchange '{exchange_name}': {e}"
+            )
+        except AMQPError as e:
+            logger.warning(f"AMQP error while deleting exchange '{exchange_name}': {e}")
+
+    @classmethod
+    async def delete_queue(
+        cls,
+        channel: AbstractChannel,
+        queue_name: str,
+        if_unused: bool = False,
+        if_empty: bool = False,
+    ) -> None:
+        """
+        Delete a queue.
+        """
+        try:
+            await channel.queue_delete(
+                queue_name=queue_name,
+                if_unused=if_unused,
+                if_empty=if_empty,
+            )
+        except ChannelNotFoundEntity:
+            # Queue might not exist, which is fine
+            logger.info(f"Queue '{queue_name}' does not exist, nothing to delete.")
+        except ChannelClosed as e:
+            logger.warning(f"Channel closed while deleting queue '{queue_name}': {e}")
+        except AMQPError as e:
+            logger.warning(f"AMQP error while deleting queue '{queue_name}': {e}")