Initial draft

src/main.py

@@ -0,0 +1,130 @@
+from fastapi import FastAPI, Depends, HTTPException, Request, Response
+from fastapi.responses import JSONResponse
+from fastapi.security.api_key import APIKeyHeader
+from starlette.exceptions import HTTPException as StarletteHTTPException
+from typing import Dict
+from pydantic import BaseModel
+from validator import verify_api_key
+from db_module import get_db, create_connection_pool, close_connection_pool, start_healthcheck_thread
+from simple_logger_handler import setup_logger, LOG_LEVEL
+from rabbitmq_handler import RabbitMQProducer
+import uvicorn
+from contextlib import asynccontextmanager
+from metrics_server import REQUEST_COUNTER
+import asyncio
+from uvicorn_logger_config import LOGGING_CONFIG
+from secret_manager import cleanup_secret_files
+
+logger = setup_logger(__name__)
+api_key_header_internal = APIKeyHeader(name="X-API-Key-Internal")
+
+class Notification(BaseModel):
+	receipent_user_id : int
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+	logger.info("Starting application...")
+
+	logger.info("Creating MySQL connection pool...")
+	create_connection_pool("ApiInternalPool")
+	start_healthcheck_thread()
+	logger.info("MySQL healthcheck thread started.")
+
+	producer = RabbitMQProducer()
+
+	try:
+		logger.info("Starting RabbitMQ producer...")
+		await producer.connect()
+		app.state.rmq_producer = producer
+		logger.info("[FastAPI] RabbitMQ producer initialized.")
+
+		SECRET_PATHS = frozenset({
+			"/etc/secrets/db_username",
+			"/etc/secrets/db_password",
+			"/etc/secrets/rmq_username",
+			"/etc/secrets/rmq_password"
+		})
+		cleanup_secret_files(SECRET_PATHS)
+
+		yield
+	finally:
+		logger.info("Closing RabbitMQ producer...")
+		await producer.close()
+
+		logger.info("Closing MySQL connection pool...")
+		close_connection_pool()
+
+api = FastAPI(
+		title="Internal Notifier API",
+		description="API to forward messages to RabbitMQ",
+		version="1.0.0",
+		lifespan=lifespan
+)
+
+@api.middleware("http")
+async def prometheus_middleware(request: Request, call_next):
+	logger.debug(f"Incoming request: {request.method} {request.url.path}")
+	status = 500
+	try:
+		response = await call_next(request)
+		status = response.status_code
+		logger.debug(f"Request processed with status: {status}")
+	except Exception as e:
+		logger.exception(f"Exception during request processing: {e}")
+		raise
+	finally:
+		REQUEST_COUNTER.labels(request.method, request.url.path, status).inc()
+		logger.debug(f"Prometheus counter incremented: {request.method} {request.url.path} {status}")
+	return response
+
+def verify_api_key_dependency_internal(db=Depends(get_db), api_key: str = Depends(api_key_header_internal)) -> str:
+	logger.debug(f"Verifying internal API key: {api_key}")
+	cursor = db.cursor()
+	cursor.execute("SELECT program_name, api_key FROM internal_api_keys WHERE status = 'active'")
+	for program_name, hashed_key in cursor.fetchall():
+		if verify_api_key(api_key=api_key, hashed=hashed_key):
+			logger.debug(f"API key verified for program: {program_name}")
+			return program_name
+	logger.warning("Internal API key verification failed")
+	raise HTTPException(status_code=403, detail="Unauthorized")
+
+@api.exception_handler(StarletteHTTPException)
+async def custom_http_exception_handler(request,exc):
+	logger.debug(f"Handling HTTP exception: {exc}")
+	if exc.status_code == 404:
+		return JSONResponse(
+			status_code=403,
+			content={"detail": "Unauthorized"}
+		)
+	return JSONResponse(
+		status_code=exc.status_code,
+		content={"detail": exc.detail}
+	)
+
+@api.get("/notifications/pull")
+async def receive_notifications(
+	notification_data: Notification,
+	request: Request,
+	db = Depends(get_db),
+	program_name: str = Depends(verify_api_key_dependency_internal)
+):
+
+	logger.info(f"Received notifcation data from {program_name} for RMQ")
+
+	messages = await request.app.state.rmq_producer.fetch_messages()
+	logger.info("Successfully retrieved messages for ")
+	print (messages)
+
+async def start_servers():
+	logger.debug("Starting FastAPI and metrics servers")
+	config_main = uvicorn.Config("main:api", host="0.0.0.0", port=8101, log_level=LOG_LEVEL.lower(), log_config=LOGGING_CONFIG)
+	config_metrics = uvicorn.Config("metrics_server:metrics_api", host="0.0.0.0", port=9000, log_level=LOG_LEVEL.lower(), log_config=LOGGING_CONFIG)
+
+	server_main = uvicorn.Server(config_main)
+	server_metrics = uvicorn.Server(config_metrics)
+
+	await asyncio.gather(server_main.serve(), server_metrics.serve())
+	logger.debug("Servers started")
+
+if __name__ == "__main__":
+	asyncio.run(start_servers())

src/rabbitmq_handler.py

@@ -0,0 +1,134 @@
+import asyncio
+import aio_pika
+from aio_pika.exceptions import AMQPException
+from secret_manager import return_credentials
+import os
+from simple_logger_handler import setup_logger
+import json
+import signal
+from urllib.parse import quote
+import weakref
+
+logger = setup_logger(__name__)
+
+class RabbitMQConsumer:
+	"""Pull-based RabbitMQ consumer. Messages are fetched on demand, one queue per user."""
+	def __init__(self):
+		self.rmq_host = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_HOST", "localhost")
+		self.rmq_vhost = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_VHOST", "app_notifications")
+		self.rmq_exchange = os.getenv("BACKEND_PUSH_NOTIFICATIONS_RMQ_EXCHANGE", "app_notifications")
+		self.pull_queue_prefix = os.getenv("BACKEND_PULL_NOTIFICATIONS_RMQ_QUEUE_PREFIX", "notifications-pull")
+		self.pull_routing_prefix = os.getenv("BACKEND_PULL_NOTIFICATIONS_RMQ_ROUTING_PREFIX", "notify.pull")
+
+		self._user_locks: weakref.WeakValueDictionary[int, asyncio.Lock] = weakref.WeakValueDictionary()
+		self.connection: aio_pika.RobustConnection | None = None
+		self.channel: aio_pika.RobustChannel | None = None
+		self.exchange: aio_pika.Exchange | None = None
+		self._closing = False
+		logger.debug(f"[RabbitMQ] Initialized consumer for exchange {self.rmq_exchange}")
+
+	def _lock_for(self, user_id: int) -> asyncio.Lock:
+		"""
+		Return the asyncio lock for this user, creating it if needed. Locks are held
+		weakly, so they are cleaned up automatically once nobody is using them.
+		"""
+		lock = self._user_locks.get(user_id)
+		if lock is None:
+			lock = asyncio.Lock()
+			self._user_locks[user_id] = lock
+		return lock
+
+	async def connect(self):
+		username = return_credentials("/etc/secrets/rmq_username")
+		password = return_credentials("/etc/secrets/rmq_password")
+		url = (
+			f"amqp://{quote(username, safe='')}:{quote(password, safe='')}"
+			f"@{self.rmq_host}/{quote(self.rmq_vhost, safe='')}"
+		)
+
+		logger.info(f"[RabbitMQ] Connecting to {self.rmq_host}...")
+		self.connection = await aio_pika.connect_robust(url)
+		self.channel = await self.connection.channel()
+		self.exchange = await self.channel.declare_exchange(
+			self.rmq_exchange, aio_pika.ExchangeType.TOPIC, durable=True
+		)
+		logger.info("[RabbitMQ] Connected to RabbitMQ")
+
+	async def ensure_user_queue(self, user_id: int) -> aio_pika.Queue:
+		"""
+		Declare the user's durable queue (with expiry and message TTL) and bind it to
+		the exchange with the user's routing key. Safe to call repeatedly, recreates it if it expired, and resets its
+		idle-expiry timer.
+		"""
+		queue = await self.channel.declare_queue(
+			f"{self.pull_queue_prefix}.{user_id}",
+			durable=True,
+			arguments={
+				"x-expires": 30 * 24 * 3600 * 1000,
+				"x-message-ttl": 7 * 24 * 3600 * 1000,
+			},
+		)
+		await queue.bind(self.exchange, routing_key=f"{self.pull_routing_prefix}.{user_id}")
+		return queue
+		
+	async def fetch_messages(self, user_id: int, max_messages: int = 1000) -> list[dict]:
+		"""Drain up to max_messages queued messages for this user and return their
+		decoded JSON bodies.
+
+		- Only one fetch per user runs at a time; different users run concurrently.
+		- Messages are acked after the whole batch has been read, so an error mid-read
+		  puts the batch back on the queue and re-raises.
+		- Undecodable messages are rejected without requeue (discarded).
+		- Returns an empty list if nothing is queued.
+		- Delivery is at-most-once towards the client: messages are acked before the
+		  HTTP response is sent.
+		"""
+
+		if self._closing:
+			raise RuntimeError("[RabbitMQ] Consumer is closing")
+		if not self.channel:
+			raise RuntimeError("[RabbitMQ] Queue not initialized, call connect() first")
+
+		user_id = int(user_id)
+
+		async with self._lock_for(user_id):  #one drain at a time
+			queue = await self.ensure_user_queue(user_id)
+			results: list[dict] = []
+			to_ack: list[aio_pika.IncomingMessage] = []
+			try:
+				while len(to_ack) < max_messages:
+					message = await queue.get(no_ack=False, fail=False, timeout=5)
+					if message is None:
+						break
+
+					try:
+						data = json.loads(message.body.decode())
+					except (json.JSONDecodeError, UnicodeDecodeError) as e:
+						logger.error(f"[RabbitMQ] Bad message, discarding: {e}", exc_info=True)
+						await message.reject(requeue=False)
+						continue
+
+					results.append(data)
+					to_ack.append(message)
+
+				for message in to_ack:
+					await message.ack()
+
+				logger.debug(f"[RabbitMQ] Fetched {len(results)} message(s)")
+				return results
+
+			except (AMQPException, OSError) as e:
+				logger.error(f"[RabbitMQ] Error while fetching, requeueing batch: {e}", exc_info=True)
+				for message in to_ack:
+					try:
+						await message.nack(requeue=True)
+					except Exception:
+						pass  # channel is gone; the broker requeues unacked messages itself
+				raise
+
+	async def close(self):
+		self._closing = True
+		if self.connection:
+			await self.connection.close()
+			logger.info("[RabbitMQ] Connection closed")
+