141 lines
4.3 KiB
Python
141 lines
4.3 KiB
Python
import asyncio
|
|
import signal
|
|
from contextlib import asynccontextmanager
|
|
|
|
import uvicorn
|
|
from aio_pika.exceptions import AMQPException
|
|
from fastapi import 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 metrics_server import REQUEST_COUNTER
|
|
from rabbitmq_handler import RabbitMQConsumer
|
|
from secret_manager import cleanup_secret_files
|
|
from simple_logger_handler import LOG_LEVEL, setup_logger
|
|
from uvicorn_logger_config import LOGGING_CONFIG
|
|
|
|
logger = setup_logger(__name__)
|
|
MAX_MESSAGES = 100
|
|
|
|
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("[FastAPI] Starting application...")
|
|
consumer = RabbitMQConsumer()
|
|
|
|
try:
|
|
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/rmq_username",
|
|
"/etc/secrets/rmq_password"
|
|
})
|
|
cleanup_secret_files(SECRET_PATHS)
|
|
|
|
yield
|
|
finally:
|
|
logger.info("[FastAPI] Closing RabbitMQ producer...")
|
|
await consumer.close()
|
|
|
|
api = FastAPI(
|
|
title="Backend Pull Notifications",
|
|
description="API to retrieve pull notifications on invocation",
|
|
version="1.0.0",
|
|
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"[FastAPI] Incoming request: {request.method} {request.url.path}")
|
|
status = 500
|
|
try:
|
|
response = await call_next(request)
|
|
status = response.status_code
|
|
logger.debug(f"[FastAPI] Request processed with status: {status}")
|
|
except Exception as e:
|
|
logger.exception(f"Exception during request processing: {e}")
|
|
raise
|
|
finally:
|
|
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
|
|
|
|
@api.exception_handler(StarletteHTTPException)
|
|
async def custom_http_exception_handler(request,exc):
|
|
logger.debug(f"[FastAPI] 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("/notifications/pull", response_model=PullResponse)
|
|
async def pull_notifications(
|
|
pull_request: PullRequest,
|
|
request: Request
|
|
):
|
|
user_id = pull_request.user_id
|
|
logger.debug(f"[FastAPI] User id {user_id} is requesting messages")
|
|
|
|
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")
|
|
|
|
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():
|
|
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)]
|
|
|
|
def stop():
|
|
for s in servers:
|
|
s.should_exit = True
|
|
|
|
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())
|