From 85d6a552c6baeecfae0834b6f36a153ef64f441e Mon Sep 17 00:00:00 2001 From: Florian Date: Sun, 11 Oct 2026 17:06:44 +0200 Subject: [PATCH] Added logic to retrieve messages via API call --- src/main.py | 119 +++++++++++++++++++++++++--------------- src/rabbitmq_handler.py | 2 +- 2 files changed, 77 insertions(+), 44 deletions(-) diff --git a/src/main.py b/src/main.py index 6ec43c7..7b197cc 100644 --- a/src/main.py +++ b/src/main.py @@ -1,42 +1,49 @@ -from fastapi import FastAPI, Depends, HTTPException, Request, Response +from fastapi import FastAPI, Depends, HTTPException, Request 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 pydantic import BaseModel, Field 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 +from rabbitmq_handler import RabbitMQConsumer import uvicorn from contextlib import asynccontextmanager +import signal from metrics_server import REQUEST_COUNTER import asyncio from uvicorn_logger_config import LOGGING_CONFIG from secret_manager import cleanup_secret_files +from aio_pika.exceptions import AMQPException logger = setup_logger(__name__) api_key_header_internal = APIKeyHeader(name="X-API-Key-Internal") +MAX_MESSAGES = 100 -class Notification(BaseModel): - receipent_user_id : int +class PullRequest(BaseModel): + user_id: int = Field(gt=0) + +class PullResponse(BaseModel): + messages: list[dict] + count: int + has_more: bool @asynccontextmanager async def lifespan(app: FastAPI): - logger.info("Starting application...") + logger.info("[FastAPI] Starting application...") - logger.info("Creating MySQL connection pool...") - create_connection_pool("ApiInternalPool") + logger.info("[FastAPI] Creating MySQL connection pool...") + create_connection_pool("ApiPullPool") start_healthcheck_thread() - logger.info("MySQL healthcheck thread started.") + logger.info("[FastAPI] MySQL healthcheck thread started.") - producer = RabbitMQProducer() + consumer = RabbitMQConsumer() try: - logger.info("Starting RabbitMQ producer...") - await producer.connect() - app.state.rmq_producer = producer - logger.info("[FastAPI] RabbitMQ producer initialized.") + logger.info("[FastAPI] Starting RabbitMQ consumer...") + await consumer.connect() + app.state.rmq_consumer = consumer + logger.info("[FastAPI] RabbitMQ consumer initialized.") SECRET_PATHS = frozenset({ "/etc/secrets/db_username", @@ -48,49 +55,56 @@ async def lifespan(app: FastAPI): yield finally: - logger.info("Closing RabbitMQ producer...") - await producer.close() + logger.info("[FastAPI] Closing RabbitMQ producer...") + await consumer.close() - logger.info("Closing MySQL connection pool...") + logger.info("[FastAPI] Closing MySQL connection pool...") close_connection_pool() api = FastAPI( - title="Internal Notifier API", - description="API to forward messages to RabbitMQ", + title="Backend Pull Notifications", + description="API to retrieve pull notifications on invocation", version="1.0.0", - lifespan=lifespan + lifespan=lifespan, + docs_url=None, + redoc_url=None, + openapi_url=None ) +KNOWN_METHODS = {"GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"} + @api.middleware("http") async def prometheus_middleware(request: Request, call_next): - logger.debug(f"Incoming request: {request.method} {request.url.path}") + logger.debug(f"[FastAPI] 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}") + logger.debug(f"[FastAPI] 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}") + route = request.scope.get("route") + path = route.path if route else "unmatched" + method = request.method if request.method in KNOWN_METHODS else "OTHER" + REQUEST_COUNTER.labels(method, path, status).inc() + logger.debug(f"[FastAPI] Prometheus counter incremented: {method} {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}") + logger.debug(f"[FastAPI] 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}") + logger.debug(f"[FastAPI] Handling HTTP exception: {exc}") if exc.status_code == 404: return JSONResponse( status_code=403, @@ -101,30 +115,49 @@ async def custom_http_exception_handler(request,exc): content={"detail": exc.detail} ) -@api.get("/notifications/pull") -async def receive_notifications( - notification_data: Notification, +@api.post("/notifications/pull", response_model=PullResponse) +async def pull_notifications( + pull_request: PullRequest, request: Request, - db = Depends(get_db), - program_name: str = Depends(verify_api_key_dependency_internal) + program_name: str = Depends(verify_api_key_dependency_internal), ): + user_id = pull_request.user_id + logger.debug(f"[FastAPI] User id {user_id} is requesting messages") - logger.info(f"Received notifcation data from {program_name} for RMQ") + try: + messages = await request.app.state.rmq_consumer.fetch_messages( + user_id=user_id, max_messages=MAX_MESSAGES + ) + except (AMQPException, OSError) as e: + logger.error(f"[FastAPI] RabbitMQ unavailable for user_id={user_id}: {e}") + raise HTTPException(status_code=503, detail="Message broker unavailable") + except RuntimeError as e: + logger.error(f"[FastAPI] Consumer not ready: {e}") + raise HTTPException(status_code=503, detail="Service not ready") - messages = await request.app.state.rmq_producer.fetch_messages() - logger.info("Successfully retrieved messages for ") - print (messages) + logger.info(f"[FastAPI] Retrieved {len(messages)} message(s) for user_id={user_id}") + return PullResponse(messages=messages, count=len(messages), has_more=len(messages) >= MAX_MESSAGES) + + +class QuietServer(uvicorn.Server): + def install_signal_handlers(self): + pass 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_main = uvicorn.Config(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) + servers = [QuietServer(config_main), QuietServer(config_metrics)] - server_main = uvicorn.Server(config_main) - server_metrics = uvicorn.Server(config_metrics) + def stop(): + for s in servers: + s.should_exit = True - await asyncio.gather(server_main.serve(), server_metrics.serve()) - logger.debug("Servers started") + loop = asyncio.get_running_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, stop) + + await asyncio.gather(*(s.serve() for s in servers)) + logger.debug("[FastAPI] Servers started") if __name__ == "__main__": asyncio.run(start_servers()) diff --git a/src/rabbitmq_handler.py b/src/rabbitmq_handler.py index d9452de..e974a62 100644 --- a/src/rabbitmq_handler.py +++ b/src/rabbitmq_handler.py @@ -71,7 +71,7 @@ class RabbitMQConsumer: 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]: + async def fetch_messages(self, user_id: int, max_messages: int = 100) -> list[dict]: """Drain up to max_messages queued messages for this user and return their decoded JSON bodies.