| | import asyncio, json, uuid |
| | import aio_pika |
| | from typing import Callable, Dict, List, Optional |
| | from .config import settings |
| |
|
| | ExchangeResolver = Callable[[str], str] |
| |
|
| | class RabbitBase: |
| | def __init__(self, exchange_type_resolver: Optional[ExchangeResolver] = None): |
| | self._conn: Optional[aio_pika.RobustConnection] = None |
| | self._chan: Optional[aio_pika.RobustChannel] = None |
| | self._exchanges: Dict[str, aio_pika.Exchange] = {} |
| | self._exchange_type_resolver = exchange_type_resolver or (lambda _: settings.RABBIT_EXCHANGE_TYPE) |
| |
|
| | async def connect(self): |
| | if self._conn and not self._conn.is_closed: |
| | return |
| | self._conn = await aio_pika.connect_robust(str(settings.AMQP_URL)) |
| | self._chan = await self._conn.channel() |
| | await self._chan.set_qos(prefetch_count=settings.RABBIT_PREFETCH) |
| |
|
| | async def ensure_exchange(self, name: str) -> aio_pika.Exchange: |
| | await self.connect() |
| | if name in self._exchanges: |
| | return self._exchanges[name] |
| | ex_type = self._exchange_type_resolver(name) |
| | ex = await self._chan.declare_exchange(name, getattr(aio_pika.ExchangeType, ex_type), durable=True) |
| | self._exchanges[name] = ex |
| | return ex |
| |
|
| | async def declare_queue_bind(self, exchange: str, queue_name: str, routing_keys: List[str], ttl_ms: Optional[int]): |
| | await self.connect() |
| | ex = await self.ensure_exchange(exchange) |
| | args = {} |
| | if ttl_ms: |
| | args["x-message-ttl"] = ttl_ms |
| | q = await self._chan.declare_queue(queue_name, durable=True, exclusive=False, auto_delete=True, arguments=args) |
| | for rk in routing_keys or [""]: |
| | await q.bind(ex, rk) |
| | return q |
| |
|
