qena-shared-lib 0.1.12__py3-none-any.whl → 0.1.13__py3-none-any.whl

qena_shared_lib/rabbitmq/_exception_handlers.py

@@ -1,8 +1,7 @@
-from typing import Annotated
+from typing import cast
 
 from pydantic import ValidationError
 
-from ..dependencies.miscellaneous import DependsOn
 from ..exceptions import (
     HTTPServiceError,
     RabbitMQServiceException,
@@ -10,138 +9,174 @@ from ..exceptions import (
     Severity,
 )
 from ..logging import LoggerProvider
-from ..logstash._base import BaseLogstashSender
+from ..remotelogging._base import BaseRemoteLogSender
 from ._listener import ListenerContext
 
 __all__ = [
-    "handle_general_mq_exception",
-    "handle_rabbit_mq_service_exception",
-    "handle_validation_error",
+    "AbstractRabbitMqExceptionHandler",
+    "GeneralMqExceptionHandler",
+    "RabbitMqServiceExceptionHandler",
+    "ValidationErrorHandler",
 ]
 
 RABBITMQ_EXCEPTION_HANDLER_LOGGER_NAME = "rabbitmq.exception_handler"
 
 
-def handle_rabbit_mq_service_exception(
-    context: ListenerContext,
-    exception: ServiceException,
-    logstash: Annotated[BaseLogstashSender, DependsOn(BaseLogstashSender)],
-    logger_provider: Annotated[LoggerProvider, DependsOn(LoggerProvider)],
-) -> None:
-    logger = logger_provider.get_logger(RABBITMQ_EXCEPTION_HANDLER_LOGGER_NAME)
-    tags = [
-        "RabbitMQ",
-        context.queue,
-        context.listener_name or "__default__",
-        exception.__class__.__name__,
-    ]
-    extra = {
-        "serviceType": "RabbitMQ",
-        "queue": context.queue,
-        "listenerName": context.listener_name,
-        "exception": exception.__class__.__name__,
-    }
-
-    match exception:
-        case HTTPServiceError() as http_service_error:
-            if http_service_error.status_code is not None:
-                str_status_code = str(http_service_error.status_code)
-                extra["statusCode"] = str_status_code
-
-                tags.append(str_status_code)
-
-            if http_service_error.response_code is not None:
-                str_response_code = str(http_service_error.response_code)
-                extra["responseCode"] = str_response_code
-
-                tags.append(str_response_code)
-        case RabbitMQServiceException() as rabbitmq_service_exception:
-            str_error_code = str(rabbitmq_service_exception.code)
-            extra["code"] = str_error_code
-
-            tags.append(str_error_code)
-
-    if exception.tags:
-        tags.extend(exception.tags)
-
-    if exception.extra:
-        extra.update(exception.extra)
-
-    exc_info = (
-        (type(exception), exception, exception.__traceback__)
-        if exception.extract_exc_info
-        else None
-    )
-
-    match exception.severity:
-        case Severity.HIGH:
-            logstash_logger_method = logstash.error
-            logger_method = logger.error
-        case Severity.MEDIUM:
-            logstash_logger_method = logstash.warning
-            logger_method = logger.warning
-        case _:
-            logstash_logger_method = logstash.info
-            logger_method = logger.info
-
-    if exception.logstash_logging:
-        logstash_logger_method(
-            message=exception.message,
-            tags=tags,
-            extra=extra,
-            exception=exception if exception.extract_exc_info else None,
-        )
-    else:
-        logger_method(
-            "\nRabbitMQ `%s` -> `%s`\n%s",
-            context.queue,
-            context.listener_name,
-            exception.message,
-            exc_info=exc_info,
-        )
+class AbstractRabbitMqExceptionHandler:
+    @property
+    def exception(self) -> type[Exception]:
+        raise NotImplementedError()
 
 
-def handle_validation_error(
-    context: ListenerContext,
-    exception: ValidationError,
-    logstash: Annotated[BaseLogstashSender, DependsOn(BaseLogstashSender)],
-) -> None:
-    logstash.error(
-        message=f"invalid rabbitmq request data at queue `{context.queue}` and listener `{context.listener_name}`",
-        tags=[
-            "RabbitMQ",
-            context.queue,
-            context.listener_name or "__default__",
-            "ValidationError",
-        ],
-        extra={
-            "serviceType": "RabbitMQ",
-            "queue": context.queue,
-            "listenerName": context.listener_name,
-            "exception": "ValidationError",
-        },
-        exception=exception,
-    )
-
-
-def handle_general_mq_exception(
-    context: ListenerContext,
-    exception: Exception,
-    logstash: Annotated[BaseLogstashSender, DependsOn(BaseLogstashSender)],
-) -> None:
-    logstash.error(
-        message=f"something went wrong while consuming message on queue `{context.queue}` and listener `{context.listener_name}`",
-        tags=[
+class RabbitMqServiceExceptionHandler(AbstractRabbitMqExceptionHandler):
+    @property
+    def exception(self) -> type[Exception]:
+        return cast(type[Exception], ServiceException)
+
+    def __init__(
+        self,
+        remote_logger: BaseRemoteLogSender,
+        logger_provider: LoggerProvider,
+    ):
+        self._logger = logger_provider.get_logger(
+            RABBITMQ_EXCEPTION_HANDLER_LOGGER_NAME
+        )
+        self._remote_logger = remote_logger
+
+    def __call__(
+        self,
+        context: ListenerContext,
+        exception: ServiceException,
+    ) -> None:
+        tags = [
             "RabbitMQ",
             context.queue,
             context.listener_name or "__default__",
             exception.__class__.__name__,
-        ],
-        extra={
+        ]
+        extra = {
             "serviceType": "RabbitMQ",
             "queue": context.queue,
             "listenerName": context.listener_name,
             "exception": exception.__class__.__name__,
-        },
-        exception=exception,
-    )
+        }
+
+        match exception:
+            case HTTPServiceError() as http_service_error:
+                if http_service_error.status_code is not None:
+                    str_status_code = str(http_service_error.status_code)
+                    extra["statusCode"] = str_status_code
+
+                    tags.append(str_status_code)
+
+                if http_service_error.response_code is not None:
+                    str_response_code = str(http_service_error.response_code)
+                    extra["responseCode"] = str_response_code
+
+                    tags.append(str_response_code)
+            case RabbitMQServiceException() as rabbitmq_service_exception:
+                str_error_code = str(rabbitmq_service_exception.code)
+                extra["code"] = str_error_code
+
+                tags.append(str_error_code)
+
+        if exception.tags:
+            tags.extend(exception.tags)
+
+        if exception.extra:
+            extra.update(exception.extra)
+
+        exc_info = (
+            (type(exception), exception, exception.__traceback__)
+            if exception.extract_exc_info
+            else None
+        )
+
+        match exception.severity:
+            case Severity.HIGH:
+                remote_logger_method = self._remote_logger.error
+                logger_method = self._logger.error
+            case Severity.MEDIUM:
+                remote_logger_method = self._remote_logger.warning
+                logger_method = self._logger.warning
+            case _:
+                remote_logger_method = self._remote_logger.info
+                logger_method = self._logger.info
+
+        if exception.remote_logging:
+            remote_logger_method(
+                message=exception.message,
+                tags=tags,
+                extra=extra,
+                exception=exception if exception.extract_exc_info else None,
+            )
+        else:
+            logger_method(
+                "\nRabbitMQ `%s` -> `%s`\n%s",
+                context.queue,
+                context.listener_name,
+                exception.message,
+                exc_info=exc_info,
+            )
+
+
+class ValidationErrorHandler(AbstractRabbitMqExceptionHandler):
+    @property
+    def exception(self) -> type[Exception]:
+        return cast(type[Exception], ValidationError)
+
+    def __init__(self, remote_logger: BaseRemoteLogSender):
+        self._remote_logger = remote_logger
+
+    def __call__(
+        self,
+        context: ListenerContext,
+        exception: ValidationError,
+    ) -> None:
+        self._remote_logger.error(
+            message=f"invalid rabbitmq request data at queue `{context.queue}` and listener `{context.listener_name}`",
+            tags=[
+                "RabbitMQ",
+                context.queue,
+                context.listener_name or "__default__",
+                "ValidationError",
+            ],
+            extra={
+                "serviceType": "RabbitMQ",
+                "queue": context.queue,
+                "listenerName": context.listener_name,
+                "exception": "ValidationError",
+            },
+            exception=exception,
+        )
+
+
+class GeneralMqExceptionHandler(AbstractRabbitMqExceptionHandler):
+    @property
+    def exception(self) -> type[Exception]:
+        return Exception
+
+    def __init__(self, remote_logger: BaseRemoteLogSender):
+        self._remote_logger = remote_logger
+
+    def __call__(
+        self,
+        context: ListenerContext,
+        exception: Exception,
+    ) -> None:
+        self._remote_logger.error(
+            message=f"something went wrong while consuming message on queue `{context.queue}` and listener `{context.listener_name}`",
+            tags=[
+                "RabbitMQ",
+                context.queue,
+                context.listener_name or "__default__",
+                exception.__class__.__name__,
+            ],
+            extra={
+                "serviceType": "RabbitMQ",
+                "queue": context.queue,
+                "listenerName": context.listener_name,
+                "exception": exception.__class__.__name__,
+            },
+            exception=exception,
+        )

qena_shared_lib/rabbitmq/_listener.py

@@ -1,12 +1,20 @@
 from abc import ABC, abstractmethod
-from asyncio import AbstractEventLoop, Future, Lock, Task, iscoroutinefunction
+from asyncio import (
+    AbstractEventLoop,
+    Future,
+    Lock,
+    Task,
+    gather,
+    iscoroutinefunction,
+)
 from dataclasses import dataclass
+from decimal import Decimal
 from functools import partial
 from inspect import Parameter, signature
 from random import uniform
 from time import time
 from types import MappingProxyType
-from typing import Any, Callable, Collection, TypeVar
+from typing import Any, Callable, Collection, TypeVar, cast
 
 from pika import BasicProperties
 from pika.adapters.asyncio_connection import AsyncioConnection
@@ -21,7 +29,7 @@ from pydantic_core import from_json, to_json
 from ..dependencies.miscellaneous import validate_annotation
 from ..exceptions import RabbitMQServiceException
 from ..logging import LoggerProvider
-from ..logstash import BaseLogstashSender
+from ..remotelogging import BaseRemoteLogSender
 from ..utils import AsyncEventLoopMixin, TypeAdapterCache
 from ._channel import BaseChannel
 from ._pool import ChannelPool
@@ -59,7 +67,7 @@ class FlowControl:
         self._channel = channel
         self._loop = loop
         self._lock = Lock()
-        self._flow_control_future: Future | None = None
+        self._flow_control_future: Future[None] | None = None
 
     async def request(self, prefetch_count: int) -> None:
         async with self._lock:
@@ -126,7 +134,7 @@ class ListenerContext:
     body: bytes
     flow_control: FlowControl
     rpc_reply: RpcReply | None = None
-    context_dispose_callback: Callable | None = None
+    context_dispose_callback: Callable[..., None] | None = None
 
     def dispose(self) -> None:
         if self.context_dispose_callback is not None:
@@ -202,7 +210,7 @@ class RetryPolicy:
 
 @dataclass
 class ListenerMethodContainer:
-    listener_method: Callable
+    listener_method: Callable[..., Any]
     parameters: MappingProxyType[str, Parameter]
     dependencies: dict[str, type]
     retry_policy: RetryPolicy | None = None
@@ -213,7 +221,7 @@ class ListenerChannelAdapter(BaseChannel):
         self,
         connection: AsyncioConnection,
         on_channel_open_callback: Callable[[Channel], None],
-        on_cancel_callback: Callable,
+        on_cancel_callback: Callable[..., None],
     ) -> None:
         super().__init__(
             connection=connection,
@@ -273,6 +281,9 @@ class Listener(AsyncEventLoopMixin):
         self._purge_on_startup = purge_on_startup
         self._retry_policy = retry_policy
         self._listeners: dict[str, ListenerMethodContainer] = {}
+        self._listeners_tasks_and_futures: list[Task[Any] | Future[Any]] = []
+        self._consumer_tag: str | None = None
+        self._cancelled = False
         self._logger = LoggerProvider.default().get_logger("rabbitmq.listener")
 
     @property
@@ -290,7 +301,7 @@ class Listener(AsyncEventLoopMixin):
     def add_listener_method(
         self,
         listener_name: str | None,
-        listener_method: Callable,
+        listener_method: Callable[..., Any],
         retry_policy: RetryPolicy | None = None,
     ) -> None:
         self._register_listener_method(
@@ -303,7 +314,7 @@ class Listener(AsyncEventLoopMixin):
     def _register_listener_method(
         self,
         listener_name: str | None,
-        listener_method: Callable,
+        listener_method: Callable[..., Any],
         parameters: MappingProxyType[str, Parameter],
         retry_policy: RetryPolicy | None = None,
     ) -> None:
@@ -343,7 +354,7 @@ class Listener(AsyncEventLoopMixin):
         channel_pool: ChannelPool,
         on_exception_callback: Callable[[ListenerContext, BaseException], bool],
         container: Container,
-        logstash: BaseLogstashSender,
+        remote_logger: BaseRemoteLogSender,
         global_retry_policy: RetryPolicy | None = None,
     ) -> None:
         self._connection = connection
@@ -351,7 +362,7 @@ class Listener(AsyncEventLoopMixin):
         self._listener_future = self.loop.create_future()
         self._on_exception_callback = on_exception_callback
         self._container = container
-        self._logstash = logstash
+        self._remote_logger = remote_logger
         self._global_retry_policy = global_retry_policy
         self._listener_channel = ListenerChannelAdapter(
             connection=connection,
@@ -362,6 +373,16 @@ class Listener(AsyncEventLoopMixin):
 
         await self._listener_future
 
+    async def cancel(self) -> None:
+        self._cancelled = True
+
+        if self._consumer_tag is not None:
+            self._channel.basic_cancel(self._consumer_tag)
+
+        _ = await gather(
+            *self._listeners_tasks_and_futures, return_exceptions=True
+        )
+
     def _on_channel_opened(self, channel: Channel) -> None:
         self._channel = channel
         self._flow_control = FlowControl(channel=self._channel, loop=self.loop)
@@ -369,6 +390,9 @@ class Listener(AsyncEventLoopMixin):
         self._declare_queue()
 
     def _on_cancelled(self) -> None:
+        if self._cancelled:
+            return
+
         self._declare_queue()
 
     def _declare_queue(self) -> None:
@@ -406,7 +430,7 @@ class Listener(AsyncEventLoopMixin):
 
     def _register_listener(self) -> None:
         try:
-            _ = self._channel.basic_consume(
+            self._consumer_tag = self._channel.basic_consume(
                 queue=self._queue,
                 auto_ack=True,
                 on_message_callback=self._on_message,
@@ -443,6 +467,9 @@ class Listener(AsyncEventLoopMixin):
         else:
             listener_name = "__default__"
 
+        if not isinstance(listener_name, str):
+            listener_name = str(listener_name)
+
         self._logger.debug(
             "message recieved from `%s` queue for listener `%s`",
             self._queue,
@@ -452,7 +479,7 @@ class Listener(AsyncEventLoopMixin):
         listener_method_container = self._listeners.get(listener_name)
 
         if listener_method_container is None:
-            self._logstash.error(
+            self._remote_logger.error(
                 message=f"no listener registered with the name `{listener_name}` on queue `{self._queue}`",
                 tags=[
                     "RabbitMQ",
@@ -485,7 +512,7 @@ class Listener(AsyncEventLoopMixin):
         )
 
     def _on_submitted_listener_error(
-        self, listener_message_meta: ListenerMessageMeta, future: Future
+        self, listener_message_meta: ListenerMessageMeta, future: Future[None]
     ) -> None:
         if future.cancelled():
             return
@@ -515,31 +542,36 @@ class Listener(AsyncEventLoopMixin):
 
             return
 
-        listener_done_callback = partial(
-            self._on_listener_done_executing, listener_message_meta
-        )
+        listener_task_or_future: Task[Any] | Future[Any] | None = None
 
         if iscoroutinefunction(
             listener_message_meta.listener_method_container.listener_method
         ):
-            self.loop.create_task(
+            listener_task_or_future = self.loop.create_task(
                 listener_message_meta.listener_method_container.listener_method(
                     *listener_method_args, **listener_method_kwargs
                 )
-            ).add_done_callback(listener_done_callback)
+            )
         else:
-            self.loop.run_in_executor(
+            listener_task_or_future = self.loop.run_in_executor(
                 executor=None,
                 func=partial(
                     listener_message_meta.listener_method_container.listener_method,
                     *listener_method_args,
                     **listener_method_kwargs,
                 ),
-            ).add_done_callback(listener_done_callback)
+            )
+
+        assert listener_task_or_future is not None
+
+        self._listeners_tasks_and_futures.append(listener_task_or_future)
+        listener_task_or_future.add_done_callback(
+            partial(self._on_listener_done_executing, listener_message_meta)
+        )
 
     def _parse_args(
         self, listener_message_meta: ListenerMessageMeta
-    ) -> tuple[list, dict]:
+    ) -> tuple[list[Any], dict[str, Any]]:
         try:
             message = from_json(listener_message_meta.body)
         except:
@@ -739,8 +771,14 @@ class Listener(AsyncEventLoopMixin):
     def _on_listener_done_executing(
         self,
         listener_message_meta: ListenerMessageMeta,
-        task_or_future: Task | Future,
+        task_or_future: Task[Any] | Future[Any],
     ) -> None:
+        if (
+            not self._cancelled
+            and task_or_future in self._listeners_tasks_and_futures
+        ):
+            self._listeners_tasks_and_futures.remove(task_or_future)
+
         if task_or_future.cancelled():
             return
 
@@ -773,11 +811,19 @@ class Listener(AsyncEventLoopMixin):
 
                     if listener_message_meta.properties.headers is not None:
                         try:
-                            times_rejected = int(
+                            _times_rejected = (
                                 listener_message_meta.properties.headers.get(
                                     "times_rejected"
                                 )
                             )
+
+                            if isinstance(_times_rejected, int):
+                                times_rejected = _times_rejected
+                            elif isinstance(
+                                _times_rejected,
+                                (str, bytes, bytearray, Decimal),
+                            ):
+                                times_rejected = int(_times_rejected)
                         except:
                             pass
 
@@ -911,7 +957,7 @@ class Listener(AsyncEventLoopMixin):
                 exception,
             )
         except:
-            self._logstash.exception(
+            self._remote_logger.exception(
                 message=f"error occured while invoking rabbitmq exception handler callback in listener `{listener_message_meta.listener_name}` and queue `{self._queue}`",
                 tags=[
                     "RabbitMQ",
@@ -930,7 +976,7 @@ class Listener(AsyncEventLoopMixin):
             return
 
         if not exception_callback_succeeded:
-            self._logstash.exception(
+            self._remote_logger.exception(
                 message=(
                     message
                     or f"error occured while handling event in listener `{listener_message_meta.listener_name}` and queue `{self._queue}`"
@@ -1013,7 +1059,9 @@ class Listener(AsyncEventLoopMixin):
         if listener_message_meta.properties.headers is None:
             listener_message_meta.properties.headers = headers
         else:
-            listener_message_meta.properties.headers.update(headers)
+            cast(
+                dict[str, Any], listener_message_meta.properties.headers
+            ).update(headers)
 
         try:
             with task.result() as channel:
@@ -1079,7 +1127,9 @@ class Listener(AsyncEventLoopMixin):
             )
         )
 
-    def _reponse_from_exception(self, exception: BaseException) -> dict:
+    def _reponse_from_exception(
+        self, exception: BaseException
+    ) -> dict[str, Any]:
         match exception:
             case RabbitMQServiceException() as rabbitmq_exception:
                 code = rabbitmq_exception.code
@@ -1123,6 +1173,8 @@ class Listener(AsyncEventLoopMixin):
             return
 
         try:
+            assert listener_message_meta.properties.reply_to is not None
+
             with task.result() as channel:
                 channel.basic_publish(
                     exchange=DEFAULT_EXCHANGE,
@@ -1168,8 +1220,8 @@ class Consumer(Listener):
 
     def consume(
         self, target: str | None = None, retry_policy: RetryPolicy | None = None
-    ) -> Callable[[Callable], Callable]:
-        def wrapper(consumer_method: Callable) -> Callable:
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        def wrapper(consumer_method: Callable[..., Any]) -> Callable[..., Any]:
             if not callable(consumer_method):
                 raise TypeError(
                     f"consumer method argument not a callable, got {type(consumer_method)}"
@@ -1201,8 +1253,8 @@ def consumer(
 
 def consume(
     target: str | None = None, retry_policy: RetryPolicy | None = None
-) -> Callable[[Callable], Callable]:
-    def wrapper(consumer_method: Callable) -> Callable:
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    def wrapper(consumer_method: Callable[..., Any]) -> Callable[..., Any]:
         if not callable(consumer_method):
             raise TypeError(
                 f"consumer method argument not a callable, got {type(consumer_method)}"
@@ -1230,8 +1282,8 @@ class RpcWorker(Listener):
 
     def execute(
         self, procedure: str | None = None
-    ) -> Callable[[Callable], Callable]:
-        def wrapper(worker_method: Callable) -> Callable:
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        def wrapper(worker_method: Callable[..., Any]) -> Callable[..., Any]:
             if not callable(worker_method):
                 raise TypeError(
                     f"worker method argument not a callable, got {type(worker_method)}"
@@ -1252,8 +1304,10 @@ def rpc_worker(queue: str, prefetch_count: int = 250) -> RpcWorker:
     return RpcWorker(queue=queue, prefetch_count=prefetch_count)
 
 
-def execute(procedure: str | None = None) -> Callable[[Callable], Callable]:
-    def wrapper(worker_method: Callable) -> Callable:
+def execute(
+    procedure: str | None = None,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    def wrapper(worker_method: Callable[..., Any]) -> Callable[..., Any]:
         if not callable(worker_method):
             raise TypeError(
                 f"worker method argument not a callable, got {type(worker_method)}"
@@ -1283,7 +1337,7 @@ class ListenerBase:
                 f"{self.__class__.__name__} not a listener, possibly no annotated with either `Consumer` or `RpcWorker`"
             )
 
-        return listener
+        return cast(Listener, listener)
 
     def register_listener_methods(self) -> Listener:
         listener = self.get_inner_listener()
@@ -1330,7 +1384,9 @@ class ListenerBase:
         attribute: Any,
         listener_method_attribute: str,
         previous_listener_method_attribute: str | None,
-    ) -> tuple[str | None, Callable | None, ListenerMethodMeta | None]:
+    ) -> tuple[
+        str | None, Callable[..., Any] | None, ListenerMethodMeta | None
+    ]:
         listener_method_meta = getattr(
             attribute, listener_method_attribute, None
         )

qena_shared_lib/rabbitmq/_rpc_client.py

@@ -27,19 +27,19 @@ __all__ = ["RpcClient"]
 
 class ExitHandler:
     _exiting = False
-    _rpc_futures: list[Future] = []
-    _original_exit_handler: Callable
+    _rpc_futures: list[Future[Any]] = []
+    _original_exit_handler: Callable[..., None]
 
     @classmethod
     def is_exising(cls) -> bool:
         return cls._exiting
 
     @classmethod
-    def add_rpc_future(cls, rpc_future: Future) -> None:
+    def add_rpc_future(cls, rpc_future: Future[Any]) -> None:
         cls._rpc_futures.append(rpc_future)
 
     @classmethod
-    def remove_rpc_future(cls, rpc_future: Future) -> None:
+    def remove_rpc_future(cls, rpc_future: Future[Any]) -> None:
         try:
             cls._rpc_futures.remove(rpc_future)
         except: