145 lines
4.8 KiB
Python
145 lines
4.8 KiB
Python
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)
|