Refactor DB and RMQ handlers: add connection pool lifecycle, remove TLS for RMQ, switch to lifespan events in FastAPI, ensure safe startup/shutdown.

src/main.py

@@ -1,15 +1,16 @@
-from fastapi import FastAPI, Depends, HTTPException
+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 validator import verify_api_key
-from db import get_db
+from db import get_db, create_connection_pool, close_connection_pool
 from logger_handler import setup_logger
-from rabbitmq_handler import send_message_to_rmq
+from rabbitmq_handler import send_message_to_rmq, create_connection, close_connection
 import uvicorn
 from uvicorn_logging_config import LOGGING_CONFIG
+from contextlib import asynccontextmanager
 
 logger = setup_logger(__name__)
 
@@ -19,10 +20,29 @@ class Notification(BaseModel):
 	receipent_user_id : int
 	message : Dict
 
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+	logger.info("Starting application...")
+
+	logger.info("Creating MySQL connection pool...")
+	create_connection_pool()
+	logger.info("Connecting to RabbitMQ...")
+	app.state.rmq_connection = create_connection()
+
+	yield
+	logger.info("Closing RabbitMQ connection...")
+	close_connection(app.state.rmq_connection)
+	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"
+		version="1.0.0",
+		lifespan=lifespan
 )
 
 def verify_api_key_dependency_internal(db=Depends(get_db), api_key: str = Depends(api_key_header_internal)) -> str:
@@ -49,15 +69,15 @@ async def custom_http_exception_handler(request,exc):
 @api.post("/internal/receive-notifications")
 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")
-	send_message_to_rmq(notification_data.user_id,notification_data.message)
+	send_message_to_rmq(request.app.state.rmq_connection,notification_data.receipent_user_id,notification_data.message)
 	logger.info("Successfully delivered message to RMQ")
 	return {"status": "queued"}
 
-
 if __name__ == "__main__":
 	uvicorn.run(
 		"main:api",