Progress!

src/rabbitmq_handler.py

@@ -1,21 +1,22 @@
 import asyncio
 import aio_pika
 from aio_pika.exceptions import AMQPException
-from secret_handler import return_credentials
+from secret_handler import return_credentials, database_lookup, decrypt_token
 import os
 from logger_handler import setup_logger
 import json
 from db import db_manager
 from send_notification import send_notification
+import signal
 
 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_PN_RMQ_HOST", "localhost:30672")
+rmq_vhost = os.getenv("BACKEND_PN_RMQ_VHOST", "app_notifications")
+rmq_exchange = os.getenv("BACKEND_PN_RMQ_EXCHANGE", "app_notifications")
 
 RABBITMQ_URL = f"amqp://{rmq_username}:{rmq_password}@{rmq_host}/{rmq_vhost}"
 
@@ -37,7 +38,7 @@ class RabbitMQConsumer:
 		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)
+		self.queue = await self.channel.declare_queue("notifications", durable=True)
 		await self.queue.bind(self.exchange, routing_key="notify.user.*")
 		logger.info("[Consumer] Connected, queue bound to notify.user.*")
 
@@ -48,7 +49,18 @@ class RabbitMQConsumer:
 			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)
+				logger.info(message.routing_key)
+				encrypted_tokens = await database_lookup(message.routing_key, db_manager)
+				if not encrypted_tokens:
+					logger.warning(f"No push tokens found for user {message.routing_key}")
+					return
+				
+				token_map = {row["uuid"]: row["token"].decode() for row in encrypted_tokens}
+				for uuid, token  in token_map.items():
+					decrypted_token = decrypt_token(token) 
+					token_map[uuid] = decrypted_token
+				await send_notification(message=data,push_tokens=token_map)
+				
 			except json.JSONDecodeError as e:
 				logger.error(f"[Consumer] Bad message, discarding: {e}")
 				await message.nack(requeue=False)
@@ -76,11 +88,17 @@ class RabbitMQConsumer:
 
 
 async def main():
+	await db_manager.connect()
 	consumer = RabbitMQConsumer(db_manager=db_manager)
 	await consumer.connect()
-	logger.info("Creating MySQL connection pool...")
-	await db_manager.connect()
+	await consumer.consume()
+	stop_event = asyncio.Event()
+	for sig in (signal.SIGINT, signal.SIGTERM):
+		asyncio.get_running_loop().add_signal_handler(sig, stop_event.set)
 
+	await stop_event.wait()
+	await consumer.close()
+	await db_manager.close()
 
 if __name__ == "__main__":
 	asyncio.run(main())