mirror of
https://github.com/ijaric/voice_assistant.git
synced 2025-05-24 14:33:26 +00:00
50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
import asyncio
|
|
import json
|
|
from contextlib import asynccontextmanager
|
|
|
|
import aio_pika
|
|
|
|
import lib.api.schemas as api_schemas
|
|
import lib.app.split_settings as app_split_settings
|
|
import lib.db.brokers as db_brokers
|
|
|
|
|
|
class RabbitMQPublisher(db_brokers.base_broker.BasePublisher):
|
|
def __init__(self, settings: app_split_settings.RabbitMQSettings()):
|
|
self.settings = settings
|
|
self.connection = None
|
|
self.channel = None
|
|
self.pool = asyncio.Queue()
|
|
self.pool_size = settings.max_pool_size
|
|
|
|
async def connect(self):
|
|
self.connection = await aio_pika.connect(self.settings.amqp_url)
|
|
self.channel = await self.connection.channel()
|
|
await self.channel.set_qos(prefetch_count=1)
|
|
exchange = await self.channel.declare_exchange(self.settings.exchange, aio_pika.ExchangeType.DIRECT)
|
|
|
|
for attr_name, attr_value in vars(self.settings.Queues).items():
|
|
if not attr_name.startswith("__"):
|
|
queue = await self.channel.declare_queue(attr_value.value)
|
|
await queue.bind(exchange, attr_value.value)
|
|
|
|
async def dispose(self):
|
|
await self.channel.close()
|
|
await self.connection.close()
|
|
|
|
async def publish_message(self, message_body: api_schemas.BrokerMessage, routing_key: str):
|
|
message = aio_pika.Message(content_type="application/json", body=json.dumps(message_body).encode())
|
|
await self.channel.default_exchange.publish(message, routing_key=routing_key)
|
|
|
|
@asynccontextmanager
|
|
async def get_connection(self):
|
|
if self.pool.empty() and self.pool.qsize() < self.pool_size:
|
|
await self.connect()
|
|
await self.pool.put(self)
|
|
|
|
conn = await self.pool.get()
|
|
try:
|
|
yield conn
|
|
finally:
|
|
await self.pool.put(conn)
|