Added logic to retrieve messages via API call

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

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.