diff --git a/src/main.py b/src/main.py index b7036fc..49df575 100644 --- a/src/main.py +++ b/src/main.py @@ -4,20 +4,17 @@ import signal from contextlib import asynccontextmanager import uvicorn -from aio_pika.exceptions import AMQPException -from fastapi import Depends, FastAPI, HTTPException, Request -from service_auth import UserId +from fastapi import Depends, FastAPI from auth_config import Caller, auth from handlers import register_handlers from metrics_server import metrics_api -from models import PullResponse from rabbitmq_handler import RabbitMQConsumer +from routers.user import user_router from secret_manager import cleanup_secret_files from simple_logger_handler import LOG_LEVEL, setup_logger logger = setup_logger(__name__) -_MAX_MESSAGES = 10 @asynccontextmanager async def lifespan(app: FastAPI): @@ -49,28 +46,8 @@ api = FastAPI( dependencies=[Depends(auth.authenticate)], ) -@api.post("/notifications/pull", response_model=PullResponse) -async def pull_notifications( - user_id: UserId, - request: Request -): - logger.debug("[FastAPI] User id %s is requesting messages", user_id) - - 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("[FastAPI] RabbitMQ unavailable for user_id=%s: %s", user_id, e) - raise HTTPException(status_code=503, detail="Message broker unavailable") from e - except RuntimeError as e: - logger.error("[FastAPI] Consumer not ready: %s", e) - raise HTTPException(status_code=503, detail="Service not ready") from e - - logger.info("[FastAPI] Retrieved %s message(s) for user_id=%s",len(messages), user_id) - return PullResponse(messages=messages, count=len(messages), has_more=len(messages) >= _MAX_MESSAGES) - register_handlers(api) +api.include_router(user_router) class QuietServer(uvicorn.Server): def install_signal_handlers(self): diff --git a/src/routers/__init__.py b/src/routers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/routers/user.py b/src/routers/user.py new file mode 100644 index 0000000..4f19c0d --- /dev/null +++ b/src/routers/user.py @@ -0,0 +1,34 @@ + +from aio_pika.exceptions import AMQPException +from fastapi import APIRouter, Depends, HTTPException, Request +from service_auth import UserId + +from auth_config import Caller, auth +from models import PullResponse +from simple_logger_handler import setup_logger + +logger = setup_logger(__name__) +_MAX_MESSAGES = 10 + +user_router = APIRouter(dependencies=[Depends(auth.require(Caller.GATEWAY))]) + +@user_router.post("/notifications/pull", response_model=PullResponse) +async def pull_notifications( + user_id: UserId, + request: Request +): + logger.debug("[FastAPI] User id %s is requesting messages", user_id) + + 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("[FastAPI] RabbitMQ unavailable for user_id=%s: %s", user_id, e) + raise HTTPException(status_code=503, detail="Message broker unavailable") from e + except RuntimeError as e: + logger.error("[FastAPI] Consumer not ready: %s", e) + raise HTTPException(status_code=503, detail="Service not ready") from e + + logger.info("[FastAPI] Retrieved %s message(s) for user_id=%s",len(messages), user_id) + return PullResponse(messages=messages, count=len(messages), has_more=len(messages) >= _MAX_MESSAGES) \ No newline at end of file