Skeleton ready

kernel/rabbitmq/client.py

@@ -0,0 +1,144 @@
+from asyncio import create_task, Event, sleep
+from time import time
+from typing import Any, Callable, Union
+from logging import Logger, INFO
+
+from aio_pika import connect_robust, Message, DeliveryMode
+
+from aio_pika.abc import AbstractIncomingMessage
+from aio_pika.abc import AbstractChannel
+from aio_pika.abc import AbstractConnection
+from aio_pika.abc import AbstractExchange
+from aio_pika.abc import AbstractQueue
+from aio_pika.exceptions import AMQPConnectionError
+
+from kernel.rabbitmq.enum.content_type import ContentType
+from kernel.rabbitmq.view.authentication_params import AuthenticationParams
+from kernel.rabbitmq.view.connection_params import ConnectionParams
+from kernel.rabbitmq.exception.connection_timeout import ConnectionTimeout
+from kernel.rabbitmq.exception.connection_not_established import ConnectionNotEstablished
+
+
+class Client:
+    __authentication_params: AuthenticationParams
+    __connection_params: ConnectionParams
+    __logger: Union[None, Logger]
+    __connection: Union[None, AbstractConnection] = None
+    __stop_event: Union[None, Event] = None
+
+    def __init__(
+        self,
+        authentication_params: AuthenticationParams,
+        connection_params: ConnectionParams,
+        logger: Union[None, Logger] = None
+    ) -> None:
+        self.__authentication_params = authentication_params
+        self.__connection_params = connection_params
+        self.__logger = logger
+
+    async def connect(self, stop_event: Event, timeout: int = 15) -> None:
+        self.__stop_event = stop_event
+
+        ran_at: float = time()
+
+        self.__log(f"Waiting for RabbitMQ connection..")
+
+        while True:
+            host = self.__connection_params.host
+            port = self.__connection_params.port
+            vhost = self.__connection_params.vhost
+
+            try:
+                self.__connection = await connect_robust(
+                    host=host,
+                    port=port,
+                    virtualhost=vhost,
+                    login=self.__authentication_params.username,
+                    password=self.__authentication_params.password,
+                )
+
+                self.__log(
+                    f"Successfully connected to {host}:{port}{vhost}"
+                )
+
+                break
+            except AMQPConnectionError:
+                seconds_wasted: float = time() - ran_at
+
+                self.__log("Connection failed! Retry after 3s..")
+
+                if seconds_wasted > timeout:
+                    raise ConnectionTimeout(
+                        f"Connection timeout to {host}:{port}{vhost}"
+                    )
+
+                await sleep(3)
+
+        create_task(self.__wait_for_stop())
+
+    async def consume(
+        self,
+        queue: str,
+        exchange: str,
+        consumer: Callable[[AbstractIncomingMessage], Any],
+        routing_key: Union[None, str] = None,
+    ) -> None:
+        self.__raise_if_disconnected()
+        assert self.__connection is not None
+
+        abstract_channel: AbstractChannel = await self.__connection.channel()
+        abstract_exchange: AbstractExchange = await abstract_channel\
+            .get_exchange(exchange)
+        abstract_queue: AbstractQueue = await abstract_channel\
+            .declare_queue(queue, durable=True)
+
+        await abstract_queue.bind(abstract_exchange, routing_key)
+
+        self.__log(
+            f"Listening queue {abstract_queue.name} binded to {abstract_exchange.name}..."
+        )
+
+        await abstract_queue.consume(consumer)
+
+    async def publish(
+        self,
+        exchange: str,
+        message: str,
+        routing_key: str = "",
+        content_type: str = ContentType.APPLICATION_JSON
+    ) -> None:
+        self.__raise_if_disconnected()
+        assert self.__connection is not None
+
+        abstract_exchange: AbstractExchange = await self.__connection.channel()\
+            .get_exchange(exchange)
+
+        amqp_message = Message(
+            body=message.encode(),
+            content_type=content_type,
+            delivery_mode=DeliveryMode.PERSISTENT,
+        )
+
+        await abstract_exchange.publish(amqp_message, routing_key=routing_key)
+
+    def __raise_if_disconnected(self) -> None:
+        if self.__connection is None or self.__connection.is_closed:
+            raise ConnectionNotEstablished("Connection not established")
+
+    async def __wait_for_stop(self) -> None:
+        if self.__stop_event is not None:
+            await self.__stop_event.wait()
+
+            self.__log(
+                "Received stop signal. Terminating RabbitMQ client.."
+            )
+
+            connection_established = self.__connection is not None \
+                and not self.__connection.is_closed
+
+            if connection_established:
+                await self.__connection.close() # type: ignore
+
+    def __log(self, message: str, level: int = INFO) -> None:
+        if self.__logger is not None:
+            self.__logger.log(level, message)