backend-api-internal/src/main.py

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())