| | |
| | from typing import Callable, Dict, List, Optional |
| | import aio_pika |
| | from urllib.parse import urlsplit, unquote |
| | from config import settings |
| |
|
| | ExchangeResolver = Callable[[str], str] |
| |
|
| |
|
| | def _parse_amqp_url(url: str) -> dict: |
| | """ |
| | Convert an AMQP URL into kwargs for aio_pika.connect_robust, |
| | so we don't pass the raw URL (and leak the password in logs). |
| | """ |
| | parts = urlsplit(url) |
| | return { |
| | "host": parts.hostname or "localhost", |
| | "port": parts.port or (5671 if parts.scheme == "amqps" else 5672), |
| | "login": parts.username or "guest", |
| | "password": parts.password or "guest", |
| | "virtualhost": unquote(parts.path[1:] or "/"), |
| | "ssl": parts.scheme == "amqps", |
| | } |
| |
|
| |
|
| | 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) -> None: |
| | if self._conn and not self._conn.is_closed: |
| | return |
| | conn_kwargs = _parse_amqp_url(str(settings.AMQP_URL)) |
| | self._conn = await aio_pika.connect_robust(**conn_kwargs) |
| | 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: Dict[str, int] = {} |
| | 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 |
| |
|
