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

qena_shared_lib/logstash/_base.py

@@ -42,11 +42,11 @@ class LogstashLogRecord:
         self._service_name = service_name
         self._log_level = log_level
         self._log_logger = log_logger
-        self._tags = None
-        self._labels = None
-        self._error_type = None
-        self._error_message = None
-        self._error_stack_trace = None
+        self._tags: list[str] | None = None
+        self._labels: dict[str, str] | None = None
+        self._error_type: str | None = None
+        self._error_message: str | None = None
+        self._error_stack_trace: str | None = None
         self._log_retries = 0
 
     @property
@@ -70,7 +70,7 @@ class LogstashLogRecord:
         return self._tags
 
     @tags.setter
-    def tags(self, tags: list[str]):
+    def tags(self, tags: list[str]) -> None:
         self._tags = tags
 
     @property
@@ -78,15 +78,14 @@ class LogstashLogRecord:
         return self._labels
 
     @labels.setter
-    def labels(self, labels: dict[str, str]):
+    def labels(self, labels: dict[str, str]) -> None:
         self._labels = labels
 
     @property
     def error(self) -> tuple[str | None, str | None, str | None]:
         return self._error_type, self._error_message, self._error_stack_trace
 
-    @error.setter
-    def error(self, exception: BaseException):
+    def error_from_exception(self, exception: BaseException) -> None:
         self._error_type = type(exception).__name__
         self._error_message = str(exception)
 
@@ -159,7 +158,7 @@ class LogstashLogRecord:
         return self._log_retries
 
     @log_retries.setter
-    def log_retries(self, log_retries: int):
+    def log_retries(self, log_retries: int) -> None:
         self._log_retries = log_retries
 
     def to_dict(self) -> dict:
@@ -222,7 +221,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         max_log_retry: int = 5,
         log_queue_size: int = 1000,
         failed_log_queue_size: int = 1000,
-    ):
+    ) -> None:
         self._sender = f"qena_shared_lib.logstash.{self.__class__.__name__}"
         self._service_name = service_name
         self._max_log_retry = max_log_retry
@@ -239,7 +238,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
             f"logstash.{self.__class__.__name__.lower()}"
         )
 
-    async def start(self):
+    async def start(self) -> None:
         if self._started:
             raise RuntimeError("logstash sender already started")
 
@@ -260,13 +259,13 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         )
         self.LOGSTASH_SENDER_STATE.state("running")
 
-    def _hook_on_start(self):
+    def _hook_on_start(self) -> None:
         pass
 
-    async def _hook_on_start_async(self):
+    async def _hook_on_start_async(self) -> None:
         pass
 
-    def _on_log_flusher_closed(self, task: Task):
+    def _on_log_flusher_closed(self, task: Task) -> None:
         if task.cancelled():
             self._close_future.set_result(None)
 
@@ -284,7 +283,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
             self._hook_on_stop_async(),
         ).add_done_callback(self._on_close_hook_done)
 
-    def _on_close_hook_done(self, task_or_future: Task | Future):
+    def _on_close_hook_done(self, task_or_future: Task | Future) -> None:
         if task_or_future.cancelled():
             self._close_future.set_result(None)
 
@@ -319,16 +318,18 @@ class BaseLogstashSender(AsyncEventLoopMixin):
 
         return self._close_future
 
-    def _on_close_future_done(self, _):
+    def _on_close_future_done(self, future: Future) -> None:
+        del future
+
         self.LOGSTASH_SENDER_STATE.state("stopped")
 
-    def _hook_on_stop(self):
+    def _hook_on_stop(self) -> None:
         pass
 
-    async def _hook_on_stop_async(self):
+    async def _hook_on_stop_async(self) -> None:
         pass
 
-    async def _flush_logs(self):
+    async def _flush_logs(self) -> None:
         while (
             not self._closed
             or not self._log_queue.empty()
@@ -400,7 +401,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
 
         raise NotImplementedError()
 
-    def _put_to_dead_letter_log_queue(self, log: LogstashLogRecord):
+    def _put_to_dead_letter_log_queue(self, log: LogstashLogRecord) -> None:
         if self._closed:
             self._logger.error(
                 "%s logger closed, falling back to stdout or stderr.\n%r",
@@ -425,7 +426,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         tags: list[str] | None = None,
         extra: dict[str, str] | None = None,
         exception: BaseException | None = None,
-    ):
+    ) -> None:
         self._enqueue_log(
             level=level,
             message=message,
@@ -440,7 +441,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         tags: list[str] | None = None,
         extra: dict[str, str] | None = None,
         exception: BaseException | None = None,
-    ):
+    ) -> None:
         self._enqueue_log(
             level=LogLevel.DEBUG,
             message=message,
@@ -455,7 +456,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         tags: list[str] | None = None,
         extra: dict[str, str] | None = None,
         exception: BaseException | None = None,
-    ):
+    ) -> None:
         self._enqueue_log(
             level=LogLevel.INFO,
             message=message,
@@ -470,7 +471,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         tags: list[str] | None = None,
         extra: dict[str, str] | None = None,
         exception: BaseException | None = None,
-    ):
+    ) -> None:
         self._enqueue_log(
             level=LogLevel.WARNING,
             message=message,
@@ -485,7 +486,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         tags: list[str] | None = None,
         extra: dict[str, str] | None = None,
         exception: BaseException | None = None,
-    ):
+    ) -> None:
         self._enqueue_log(
             level=LogLevel.ERROR,
             message=message,
@@ -500,7 +501,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         tags: list[str] | None = None,
         extra: dict[str, str] | None = None,
         exception: BaseException | None = None,
-    ):
+    ) -> None:
         if exception is None:
             _, exception, _ = exc_info()
 
@@ -511,7 +512,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
             exception=exception,
         )
 
-    def set_level(self, level: LogLevel):
+    def set_level(self, level: LogLevel) -> None:
         self._level = level
 
     def _enqueue_log(
@@ -521,7 +522,7 @@ class BaseLogstashSender(AsyncEventLoopMixin):
         tags: list[str] | None = None,
         extra: dict[str, str] | None = None,
         exception: BaseException | None = None,
-    ):
+    ) -> None:
         if self._closed:
             self._logger.warning("Logstash http logger is already close")
 
@@ -568,6 +569,6 @@ class BaseLogstashSender(AsyncEventLoopMixin):
             log.labels = extra
 
         if exception:
-            log.error = exception
+            log.error_from_exception(exception)
 
         return log

qena_shared_lib/logstash/_http_sender.py

@@ -57,5 +57,5 @@ class HTTPSender(BaseLogstashSender):
 
         return SenderResponse(sent=True)
 
-    async def _hook_on_stop_async(self):
+    async def _hook_on_stop_async(self) -> None:
         await self._client.aclose()

qena_shared_lib/logstash/_tcp_sender.py

@@ -1,4 +1,4 @@
-from asyncio import open_connection
+from asyncio import StreamWriter, open_connection
 
 from pydantic_core import to_json
 
@@ -33,21 +33,21 @@ class TCPSender(BaseLogstashSender):
 
         return SenderResponse(sent=True)
 
-    async def _hook_on_start_async(self):
+    async def _hook_on_start_async(self) -> None:
         await self._client.open()
 
-    async def _hook_on_stop_async(self):
+    async def _hook_on_stop_async(self) -> None:
         await self._client.close()
 
 
 class AsyncTcpClient:
-    def __init__(self, host: str, port: int):
+    def __init__(self, host: str, port: int) -> None:
         self._host = host
         self._port = port
-        self._writer = None
+        self._writer: StreamWriter | None = None
         self._client_closed = False
 
-    async def write(self, json: dict):
+    async def write(self, json: dict) -> None:
         if self._client_closed:
             raise RuntimeError("async tcp client already closed")
 
@@ -64,10 +64,10 @@ class AsyncTcpClient:
         except ConnectionResetError:
             await self._connect()
 
-    async def open(self):
+    async def open(self) -> None:
         await self._connect()
 
-    async def close(self):
+    async def close(self) -> None:
         if not self._writer_closed() and self._writer is not None:
             self._writer.write_eof()
             await self._writer.drain()
@@ -78,7 +78,7 @@ class AsyncTcpClient:
     def _writer_closed(self) -> bool:
         return self._writer is None or self._writer.is_closing()
 
-    async def _connect(self):
+    async def _connect(self) -> None:
         _, self._writer = await open_connection(
             host=self._host, port=self._port
         )

qena_shared_lib/rabbitmq/__init__.py

@@ -22,6 +22,7 @@ from ._listener import (
     execute,
     rpc_worker,
 )
+from ._pool import ChannelPool
 from ._publisher import Publisher
 from ._rpc_client import RpcClient
 
@@ -29,6 +30,7 @@ __all__ = [
     "AbstractRabbitMQService",
     "BackoffRetryDelay",
     "BaseChannel",
+    "ChannelPool",
     "consume",
     "CONSUMER_ATTRIBUTE",
     "consumer",

qena_shared_lib/rabbitmq/_base.py

@@ -10,6 +10,7 @@ from functools import partial
 from inspect import Parameter, signature
 from random import uniform
 from typing import (
+    Any,
     Awaitable,
     Callable,
     Concatenate,
@@ -61,6 +62,7 @@ ExceptionHandler = (
     Callable[Concatenate[ListenerContext, E, P], None]
     | Callable[Concatenate[ListenerContext, E, P], Awaitable]
 )
+R = TypeVar("R")
 
 
 class AbstractRabbitMQService(ABC):
@@ -113,11 +115,11 @@ class RabbitMqManager(AsyncEventLoopMixin):
         self._reconnect_delay_jitter = reconnect_delay_jitter
         self._container = container or Container()
         self._listener_global_retry_policy = listener_global_retry_policy
-        self._connection = None
+        self._connection: AsyncioConnection | None = None
         self._connected = False
         self._disconnected = False
         self._connection_blocked = False
-        self._services = []
+        self._services: list[AbstractRabbitMQService] = []
         self._exception_handlers: dict[
             type[Exception], ExceptionHandlerContainer
         ] = {}
@@ -146,7 +148,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
         def wrapper(
             handler: ExceptionHandler,
-        ):
+        ) -> ExceptionHandler:
             if not callable(handler):
                 raise TypeError(f"handler not a callable, got {type(handler)}")
 
@@ -201,7 +203,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
         return True, None
 
-    def include_listener(self, listener: Listener | type[ListenerBase]):
+    def include_listener(self, listener: Listener | type[ListenerBase]) -> None:
         if isinstance(listener, Listener):
             self._listeners.append(listener)
 
@@ -216,7 +218,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
             f"listener is {type(listener)}, expected instance of type or subclass of `Listener` or `type[ListenerBase]`"
         )
 
-    def _register_listener_classes(self, listener_class: type):
+    def _register_listener_classes(self, listener_class: type) -> None:
         inner_listener = getattr(listener_class, LISTENER_ATTRIBUTE, None)
 
         if inner_listener is None:
@@ -239,7 +241,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
         self,
         rabbit_mq_service: AbstractRabbitMQService
         | type[AbstractRabbitMQService],
-    ):
+    ) -> None:
         if not isinstance(rabbit_mq_service, AbstractRabbitMQService) and (
             not isinstance(rabbit_mq_service, type)
             or not issubclass(rabbit_mq_service, AbstractRabbitMQService)
@@ -276,13 +278,13 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
         return self._connected_future
 
-    def _resolve_listener_classes(self):
+    def _resolve_listener_classes(self) -> None:
         self._listeners.extend(
             listener.register_listener_methods()
             for listener in self._container.resolve_all(ListenerBase)
         )
 
-    def _resolve_service_classes(self):
+    def _resolve_service_classes(self) -> None:
         self._services.extend(
             self._container.resolve_all(AbstractRabbitMQService)
         )
@@ -323,9 +325,9 @@ class RabbitMqManager(AsyncEventLoopMixin):
         exchange: str | None = None,
         procedure: str | None = None,
         headers: dict[str, str] | None = None,
-        return_type: type | None = None,
-        timeout: float = 0,
-    ) -> RpcClient:
+        return_type: type[R] | None = None,
+        timeout: float = 15,
+    ) -> RpcClient[R]:
         if timeout == 0:
             self._logger.warning(
                 "rpc call with 0 seconds timeout may never return back"
@@ -355,7 +357,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
             and not self._connection.is_closed
         )
 
-    def disconnect(self):
+    def disconnect(self) -> None:
         if self._disconnected:
             raise RuntimeError("already disconnected from rabbitmq")
 
@@ -372,7 +374,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
         self.RABBITMQ_CONNECTION_STATE.state("disconnected")
 
-    def _on_connection_opened(self, connection: AsyncioConnection):
+    def _on_connection_opened(self, connection: AsyncioConnection) -> None:
         self._connection = connection
         self._connection_blocked = False
 
@@ -394,7 +396,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
             self._channel_pool.fill(self._connection)
         ).add_done_callback(self._channel_pool_filled)
 
-    def _channel_pool_drained(self, task: Task):
+    def _channel_pool_drained(self, task: Task) -> None:
         if task.cancelled():
             if not self._connected and not self._connected_future.done():
                 _ = self._connected_future.cancel(None)
@@ -422,7 +424,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
     def _on_connection_blocked(
         self, connection: AsyncioConnection, method: Method
-    ):
+    ) -> None:
         del connection, method
 
         self._connection_blocked = True
@@ -434,7 +436,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
     def _on_connection_unblocked(
         self, connection: AsyncioConnection, method: Method
-    ):
+    ) -> None:
         del connection, method
 
         self._connection_blocked = False
@@ -445,7 +447,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
     def _is_connection_blocked(self) -> bool:
         return self._connection_blocked
 
-    def _channel_pool_filled(self, task: Task):
+    def _channel_pool_filled(self, task: Task) -> None:
         if task.cancelled():
             if not self._connected and not self._connected_future.done():
                 _ = self._connected_future.cancel(None)
@@ -502,7 +504,9 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
             return [listener_configuration_error_future]
 
-    def _initialize_services(self) -> list[Awaitable]:
+    def _initialize_services(self) -> list[Future]:
+        assert self._connection is not None
+
         try:
             return [
                 service.initialize(self._connection, self._channel_pool)
@@ -515,7 +519,9 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
             return [service_initialization_error_future]
 
-    def _listener_and_service_config_and_init_done(self, future: Future):
+    def _listener_and_service_config_and_init_done(
+        self, future: Future
+    ) -> None:
         if future.cancelled():
             if not self._connected and not self._connected_future.done():
                 _ = self._connected_future.cancel(None)
@@ -564,7 +570,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
     def _on_connection_open_error(
         self, connection: AsyncioConnection, exception: BaseException
-    ):
+    ) -> None:
         del connection
 
         if not self._connected and not self._connected_future.done():
@@ -629,8 +635,8 @@ class RabbitMqManager(AsyncEventLoopMixin):
         return True
 
     def _on_exception_handler_done(
-        self, context: ListenerContext, task_or_future: Task | Future
-    ):
+        self, context: ListenerContext, task_or_future: Task[Any] | Future[Any]
+    ) -> None:
         if task_or_future.cancelled():
             return
 
@@ -646,7 +652,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
     def _on_connection_closed(
         self, connection: AsyncioConnection, exception: BaseException
-    ):
+    ) -> None:
         del connection
 
         if not self._connected and not self._connected_future.done():
@@ -661,7 +667,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
             )
             self._reconnect()
 
-    def _reconnect(self):
+    def _reconnect(self) -> None:
         self.RABBITMQ_CONNECTION_STATE.state("reconnecting")
         self.loop.call_later(
             delay=self._reconnect_delay
@@ -669,7 +675,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
             callback=self._on_time_to_reconnect,
         )
 
-    def _on_time_to_reconnect(self):
+    def _on_time_to_reconnect(self) -> None:
         try:
             connected_future = self.connect()
         except Exception as e:
@@ -689,7 +695,7 @@ class RabbitMqManager(AsyncEventLoopMixin):
 
         connected_future.add_done_callback(self._on_reconnect_done)
 
-    def _on_reconnect_done(self, future: Future):
+    def _on_reconnect_done(self, future: Future[None]) -> None:
         if future.cancelled():
             return
 

qena_shared_lib/rabbitmq/_channel.py

@@ -1,4 +1,5 @@
 from asyncio import Future
+from inspect import Traceback
 from random import uniform
 from uuid import UUID, uuid4
 
@@ -26,7 +27,7 @@ class BaseChannel(AsyncEventLoopMixin):
         self._reopen_delay_jitter = reopen_delay_jitter
         self._failed_reopen_threshold = failed_reopen_threshold
         self._completion_future = self.loop.create_future()
-        self._channel = None
+        self._channel: Channel | None = None
         self._channel_id = uuid4()
         self._reopen_failures = 0
         self._reserved = False
@@ -51,13 +52,13 @@ class BaseChannel(AsyncEventLoopMixin):
 
         return self._completion_future
 
-    def reserve(self):
+    def reserve(self) -> None:
         self._reserved = True
 
-    def release(self):
+    def release(self) -> None:
         self._reserved = False
 
-    def _on_channel_opened(self, channel: Channel):
+    def _on_channel_opened(self, channel: Channel) -> None:
         self._channel = channel
 
         try:
@@ -66,7 +67,7 @@ class BaseChannel(AsyncEventLoopMixin):
             if not self._completion_future.done():
                 self._completion_future.set_exception(e)
             else:
-                if not self.channel_closed():
+                if self._channel is not None and not self.channel_closed():
                     self._channel.close()
 
                 self._reopen()
@@ -78,7 +79,7 @@ class BaseChannel(AsyncEventLoopMixin):
 
         self._completion_future.set_result(self._channel_id)
 
-    def _hook_on_channel_opened(self):
+    def _hook_on_channel_opened(self) -> None:
         pass
 
     def channel_closed(self) -> bool:
@@ -87,7 +88,7 @@ class BaseChannel(AsyncEventLoopMixin):
 
         return self._channel.is_closing or self._channel.is_closed
 
-    def _on_cancelled(self, method: Basic.Cancel):
+    def _on_cancelled(self, method: Basic.Cancel) -> None:
         del method
 
         if not self._completion_future.done():
@@ -102,10 +103,12 @@ class BaseChannel(AsyncEventLoopMixin):
         except:
             self._logger.exception("error occured on invoking cancel hook")
 
-    def _hook_on_cancelled(self):
+    def _hook_on_cancelled(self) -> None:
         pass
 
-    def _on_channel_closed(self, channel: Channel, error: BaseException):
+    def _on_channel_closed(
+        self, channel: Channel, error: BaseException
+    ) -> None:
         del channel
 
         try:
@@ -118,10 +121,10 @@ class BaseChannel(AsyncEventLoopMixin):
         if not isinstance(error, ChannelClosedByClient):
             self._reopen()
 
-    def _hook_on_channel_closed(self, error: BaseException):
+    def _hook_on_channel_closed(self, error: BaseException) -> None:
         del error
 
-    def _reopen(self):
+    def _reopen(self) -> None:
         if (
             self._failed_reopen_threshold is not None
             and self._reopen_failures >= self._failed_reopen_threshold
@@ -135,7 +138,7 @@ class BaseChannel(AsyncEventLoopMixin):
             callback=self._on_time_to_reopen,
         )
 
-    def _on_time_to_reopen(self):
+    def _on_time_to_reopen(self) -> None:
         if self._connection.is_closing or self._connection.is_closed:
             self._can_be_disposed = True
 
@@ -151,7 +154,7 @@ class BaseChannel(AsyncEventLoopMixin):
             )
             self._reopen()
 
-    def __enter__(self):
+    def __enter__(self) -> Channel:
         if not self._reserved:
             self.reserve()
 
@@ -160,7 +163,14 @@ class BaseChannel(AsyncEventLoopMixin):
 
         return self._channel
 
-    def __exit__(self, *_):
+    def __exit__(
+        self,
+        exception_type: type[Exception],
+        exception_value: Exception,
+        exception_traceback: Traceback,
+    ) -> None:
+        del exception_type, exception_value, exception_traceback
+
         self.release()
 
     @property

qena_shared_lib/rabbitmq/_exception_handlers.py

@@ -27,7 +27,7 @@ def handle_rabbit_mq_service_exception(
     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",
@@ -105,7 +105,7 @@ 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=[
@@ -128,7 +128,7 @@ 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=[