First commit

src/rabbitmq_handler.py

@@ -0,0 +1,87 @@
+import asyncio
+import aio_pika
+from aio_pika.exceptions import AMQPException
+from secret_handler import return_credentials
+import os
+from logger_handler import setup_logger
+import json
+from db import db_manager
+from send_notification import send_notification
+
+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")
+
+RABBITMQ_URL = f"amqp://{rmq_username}:{rmq_password}@{rmq_host}/{rmq_vhost}"
+
+
+class RabbitMQConsumer:
+	def __init__(self, url=RABBITMQ_URL, db_manager=db_manager, 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: aio_pika.Queue | None = None
+		self._closing = False
+		self.db_manager = db_manager
+
+	async def connect(self):
+		self.connection = await aio_pika.connect_robust(self.url)
+		self.channel = await self.connection.channel()
+		self.exchange = await self.channel.declare_exchange(
+			self.exchange_name, aio_pika.ExchangeType.TOPIC, durable=True
+		)
+		self.queue = await self.channel.declare_queue("backend_push_notifications", durable=True)
+		await self.queue.bind(self.exchange, routing_key="notify.user.*")
+		logger.info("[Consumer] Connected, queue bound to notify.user.*")
+
+	async def handle_message(self, message: aio_pika.IncomingMessage):
+		if self._closing:
+			return
+		async with message.process():
+			try:
+				data = json.loads(message.body.decode())
+				logger.info(f"[Consumer] Received: {data}")
+				await send_notification(routing_key=message.routing_key,message=data,db_manager=self.db_manager)
+			except json.JSONDecodeError as e:
+				logger.error(f"[Consumer] Bad message, discarding: {e}")
+				await message.nack(requeue=False)
+			except AMQPException as e:
+				logger.error(f"[Consumer] RabbitMQ error: {e}")
+				await message.nack(requeue=True)
+			except Exception as e:
+				logger.critical(f"[Consumer] Fatal error: {e}")
+				raise
+
+	async def consume(self):
+		if self._closing:
+			return
+		if not self.queue:
+			raise RuntimeError("Queue not initialized")
+		await self.queue.consume(self.handle_message, no_ack=False)
+		logger.info("[Consumer] Consuming messages...")
+
+	async def close(self):
+		self._closing = True
+		if self.connection:
+			await self.connection.close()
+			logger.info("[aio-pika] Connection closed.")
+
+
+
+async def main():
+	consumer = RabbitMQConsumer(db_manager=db_manager)
+	await consumer.connect()
+	logger.info("Creating MySQL connection pool...")
+	await db_manager.connect()
+
+
+if __name__ == "__main__":
+	asyncio.run(main())
+