From 2bbb9dff3c23cff938fe3cdc1b872d9094daf843 Mon Sep 17 00:00:00 2001 From: Florian Date: Sun, 11 Oct 2026 12:39:05 +0200 Subject: [PATCH] Initial draft --- src/main.py | 130 ++++++++++++++++++++++++++++++++++++++ src/rabbitmq_handler.py | 134 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 264 insertions(+) create mode 100644 src/main.py create mode 100644 src/rabbitmq_handler.py diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..6ec43c7 --- /dev/null +++ b/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()) diff --git a/src/rabbitmq_handler.py b/src/rabbitmq_handler.py new file mode 100644 index 0000000..d9452de --- /dev/null +++ b/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") +