import asyncio import os from contextlib import asynccontextmanager from typing import Dict import httpx import uvicorn from fastapi import Depends, FastAPI, HTTPException, Request from fastapi.responses import JSONResponse from fastapi.security.api_key import APIKeyHeader from pydantic import BaseModel, Field from starlette.exceptions import HTTPException as StarletteHTTPException from db_module import close_connection_pool, create_connection_pool, get_db, start_healthcheck_thread from metrics_server import REQUEST_COUNTER from rabbitmq_handler import RabbitMQProducer from secret_manager import cleanup_secret_files from simple_logger_handler import LOG_LEVEL, setup_logger from validator import verify_api_key logger = setup_logger(__name__) api_key_header_internal = APIKeyHeader(name="X-API-Key-Internal") class Notification(BaseModel): user_id : int = Field(gt=0) message : Dict class PullResponse(BaseModel): messages: list[dict] count: int has_more: bool class PullRequest(BaseModel): user_id: int = Field(gt=0) @asynccontextmanager async def lifespan(app: FastAPI): logger.info("[FastAPI] Starting application...") logger.info("[FastAPI] Creating MySQL connection pool...") create_connection_pool("ApiInternalPool") start_healthcheck_thread() logger.info("[FastAPI] MySQL healthcheck thread started.") producer = RabbitMQProducer() PULL_SERVICE_URL = os.getenv( "BACKEND_PULL_NOTIFICATIONS_URL", "http://backend-pull-notifications.app-notifications.svc.cluster.local:8101", ) try: logger.info("[FastAPI] Starting RabbitMQ producer...") await producer.connect() app.state.rmq_producer = producer logger.info("[FastAPI] RabbitMQ producer initialized.") app.state.http_client = httpx.AsyncClient ( base_url=PULL_SERVICE_URL, timeout=httpx.Timeout(10.0,connect=3.0) ) logger.info("[FastAPI] HTTP client 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("[FastAPI] Closing RabbitMQ producer...") await producer.close() logger.info("[FastAPI] Closing MySQL connection pool...") close_connection_pool() await app.state.http_client.aclose() logger.info("[FastAPI] Closing Http client") 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"[FastAPI] Exception during request processing: {e}") raise finally: REQUEST_COUNTER.labels(request.method, request.url.path, status).inc() logger.debug(f"[FastAPI] 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: with db.cursor() as cursor: cursor.execute("SELECT program_name, api_key FROM internal_api_keys WHERE status = 'active'") rows = cursor.fetchall() for program_name, hashed_key in rows: if verify_api_key(api_key=api_key, hashed=hashed_key): logger.debug(f"[FastAPI] API key verified for program: {program_name}") return program_name logger.warning("[FastAPI] 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.post("/internal/receive-notifications") 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"[FastAPI] Received notifcation data from {program_name} for RMQ") logger.debug(f"[FastAPI] Notification data: {notification_data}") await request.app.state.rmq_producer.publish( notification_data.user_id, notification_data.message ) logger.info("Successfully queued for delivery to RMQ") return {"status": "queued"} @api.post("/internal/pull-notifications", response_model=PullResponse) async def pull_notifications( pull_request: PullRequest, request: Request, 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") client: httpx.AsyncClient = request.app.state.http_client try: resp = await client.post("/internal/pull-notifications", json={"user_id": user_id}) resp.raise_for_status() data = PullResponse.model_validate(resp.json()) except (httpx.HTTPError, ValueError) as e: logger.error(f"[FastAPI] Pull service failed for user_id={user_id}: {e!r}") status = 504 if isinstance(e, httpx.TimeoutException) else 502 raise HTTPException(status_code=status, detail="Notification service unavailable") from e logger.info(f"[FastAPI] Returning {len(data.messages)} message(s) for user_id={user_id}, has_more={data.has_more}") return data async def start_servers(): logger.info("[Server] Starting main API (port 8101) and metrics server (port 9000)...") config_main = uvicorn.Config("main:api", host="0.0.0.0", port=8101, log_level=LOG_LEVEL.lower(), log_config=None) config_metrics = uvicorn.Config("metrics_server:metrics_api", host="0.0.0.0", port=9000, log_level=LOG_LEVEL.lower(), log_config=None) server_main = uvicorn.Server(config_main) server_metrics = uvicorn.Server(config_metrics) await asyncio.gather(server_main.serve(), server_metrics.serve()) logger.info("[Server] Both servers have stopped.") if __name__ == "__main__": asyncio.run(start_servers())