Improved RabbitMQ handling
- Switched to async because it offers easy automatic reconnect features on connection failure - Adjusted notification path to reflect that change - Added in memory message queue if RMQ is not reachable
This commit is contained in:
florian committed
1 parent
c0fcdaeb4f
commit
873e7e7854
5 files changed
+104
-120
No files matched your search
+74
-60
@@ -1,72 +1,86 @@
|
||||
import pika
|
||||
from typing import Dict
|
||||
import asyncio
|
||||
import aio_pika
|
||||
from secret_handler import return_credentials
|
||||
import json
|
||||
import time
|
||||
import sys
|
||||
import os
|
||||
from logger_handler import setup_logger
|
||||
import json
|
||||
|
||||
logger = setup_logger(__name__)
|
||||
|
||||
rmq_username = return_credentials("/etc/secrets/rmq_username")
|
||||
rmq_password = return_credentials("/etc/secrets/rmq_password")
|
||||
rmq_host = os.getenv("BACKEND_API_INTERNAL_RMQ_HOST","localhost")
|
||||
rmq_vhost = os.getenv("BACKEND_API_INTERNAL_RMQ_VHOST","app_notifications")
|
||||
rmq_exchange = os.getenv("BACKEND_API_INTERNAL_RMQ_EXCHANGE","app_notifications")
|
||||
rmq_host = os.getenv("BACKEND_API_INTERNAL_RMQ_HOST", "localhost")
|
||||
rmq_vhost = os.getenv("BACKEND_API_INTERNAL_RMQ_VHOST", "app_notifications")
|
||||
rmq_exchange = os.getenv("BACKEND_API_INTERNAL_RMQ_EXCHANGE", "app_notifications")
|
||||
|
||||
MAX_RETRIES = 5
|
||||
RETRY_DELAY = 5
|
||||
RABBITMQ_URL = f"amqp://{rmq_username}:{rmq_password}@{rmq_host}/{rmq_vhost}"
|
||||
|
||||
|
||||
def create_connection():
|
||||
credentials = pika.PlainCredentials(username=rmq_username, password=rmq_password)
|
||||
conn_params = pika.ConnectionParameters(
|
||||
host=rmq_host,
|
||||
port=5672,
|
||||
credentials=credentials,
|
||||
virtual_host=rmq_vhost
|
||||
)
|
||||
class RabbitMQProducer:
|
||||
def __init__(self, url=RABBITMQ_URL, exchange_name=rmq_exchange):
|
||||
self.url = url
|
||||
self.exchange_name = exchange_name
|
||||
self.connection: aio_pika.RobustConnection | None = None
|
||||
self.channel: aio_pika.RobustChannel | None = None
|
||||
self.exchange: aio_pika.Exchange | None = None
|
||||
self._queue: asyncio.Queue[tuple[int, dict]] = asyncio.Queue()
|
||||
self._flush_task: asyncio.Task | None = None
|
||||
self._closing = False
|
||||
self._ready = asyncio.Event()
|
||||
|
||||
for attempt in range(1, MAX_RETRIES + 1):
|
||||
try:
|
||||
connection = pika.BlockingConnection(conn_params)
|
||||
print("[RMQ] Connection established.")
|
||||
return connection
|
||||
except Exception as e:
|
||||
print(f"[RMQ] Attempt {attempt} failed: {e}")
|
||||
if attempt < MAX_RETRIES:
|
||||
time.sleep(RETRY_DELAY)
|
||||
else:
|
||||
print("[RMQ] Failed to connect after maximum retries — exiting.")
|
||||
sys.exit(1)
|
||||
async def connect(self):
|
||||
self.connection = await aio_pika.connect_robust(self.url)
|
||||
self.channel = await self.connection.channel(publisher_confirms=True)
|
||||
self.exchange = await self.channel.declare_exchange(
|
||||
self.exchange_name, aio_pika.ExchangeType.TOPIC, durable=True
|
||||
)
|
||||
logger.info(f"[aio-pika] Connected and exchange '{self.exchange_name}' ready.")
|
||||
self._ready.set()
|
||||
self._flush_task = asyncio.create_task(self._flush_queue_loop())
|
||||
|
||||
async def publish(self, routing_key: int, message: dict):
|
||||
await self._queue.put((routing_key, message))
|
||||
logger.debug(f"[RabbitMQ] Queued message for {routing_key}")
|
||||
|
||||
async def _flush_queue_loop(self):
|
||||
await self._ready.wait()
|
||||
logger.debug(f"here")
|
||||
while True:
|
||||
routing_key, message = await self._queue.get()
|
||||
try:
|
||||
await self.exchange.publish(
|
||||
aio_pika.Message(
|
||||
body=json.dumps(message).encode(),
|
||||
content_type="application/json",
|
||||
delivery_mode=aio_pika.DeliveryMode.PERSISTENT
|
||||
),
|
||||
routing_key=f"notify.user.{routing_key}",
|
||||
mandatory=True
|
||||
)
|
||||
logger.debug(f"[aio-pika] Published message to notify.user.{routing_key}")
|
||||
except Exception as e:
|
||||
logger.warning(f"[aio-pika] Publish failed, requeuing: {e}")
|
||||
await asyncio.sleep(2)
|
||||
await self._queue.put((routing_key, message))
|
||||
finally:
|
||||
self._queue.task_done()
|
||||
|
||||
async def close(self):
|
||||
self._closing = True
|
||||
if self._flush_task and not self._flush_task.done():
|
||||
self._flush_task.cancel()
|
||||
try:
|
||||
await self._flush_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
if self.connection:
|
||||
await self.connection.close()
|
||||
logger.info("[aio-pika] Connection closed.")
|
||||
|
||||
|
||||
def send_message_to_rmq(connection, user_id: int, message: Dict):
|
||||
if not connection or not connection.is_open:
|
||||
print("[RMQ] Connection lost, reconnecting...")
|
||||
connection = create_connection()
|
||||
async def main():
|
||||
producer = RabbitMQProducer()
|
||||
await producer.connect()
|
||||
await producer.publish("1", {"type": "notification", "content": "Test message"})
|
||||
await producer.close()
|
||||
|
||||
channel = connection.channel()
|
||||
channel.exchange_declare(exchange=rmq_exchange, exchange_type="topic", durable=True)
|
||||
channel.confirm_delivery()
|
||||
|
||||
channel.basic_publish(
|
||||
exchange=rmq_exchange,
|
||||
routing_key=f"notify.user.{user_id}",
|
||||
body=json.dumps(message),
|
||||
properties=pika.BasicProperties(
|
||||
content_type="application/json",
|
||||
delivery_mode=2
|
||||
),
|
||||
mandatory=True
|
||||
)
|
||||
|
||||
|
||||
def close_connection(connection):
|
||||
if connection and connection.is_open:
|
||||
connection.close()
|
||||
print("[RMQ] Connection closed.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
conn = create_connection()
|
||||
send_message_to_rmq(conn, 1, {"type": "notification", "content": "Vault TLS cert reloaded successfully."})
|
||||
close_connection(conn)
|
||||
Reference in new issue
Block a user