pythonProject2/main.py - lab4 - Pepegator/task-api

pythonProject2/main.py
117 строк3 KB
Pol136
add rabbitmq
30 дек 2025, 17:44
30 дек 2025, 17:447c8d717
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from slowapi import Limiter
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
from slowapi.middleware import SlowAPIMiddleware
from contextlib import asynccontextmanager
import logging
import asyncio

from database import engine, Base
from routers.v1 import users as users_v1, tasks as tasks_v1
from routers.v2 import users as users_v2, tasks as tasks_v2
from routers.internal import tasks as internal_tasks, stats as internal_stats
from rabbitmq_service import start_consumer, stop_consumer

Base.metadata.create_all(bind=engine)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

limiter = Limiter(
    key_func=get_remote_address,
    default_limits=["60/minute"]
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Task Management API запущен")

    consumer_task = None
    try:
        consumer_task = asyncio.create_task(start_consumer())
        logger.info("✓ RabbitMQ Consumer запущен успешно")
    except Exception as e:
        logger.error(f"✗ Ошибка при запуске RabbitMQ Consumer: {e}")
        logger.warning("Продолжаю работу - REST API будет доступен, RabbitMQ недоступен")
        consumer_task = None

    yield

    if consumer_task and not consumer_task.done():
        consumer_task.cancel()
        try:
            await consumer_task
        except asyncio.CancelledError:
            pass

    await stop_consumer()
    logger.info("Приложение остановлено")


app = FastAPI(
    title="Task Management API",
    description="REST API для управления задачами",
    version="1.0.0",
    lifespan=lifespan
)


app.state.limiter = limiter
app.add_middleware(SlowAPIMiddleware)


@app.exception_handler(RateLimitExceeded)
async def rate_limit_handler(request: Request, exc: RateLimitExceeded):
    return {
        "detail": "Rate limit exceeded",
        "retry_after": 60
    }


app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.middleware("http")
async def add_rate_limit_headers(request: Request, call_next):
    response = await call_next(request)

    response.headers["X-RateLimit-Limit"] = "60"
    response.headers["X-RateLimit-Remaining"] = "59"
    response.headers["X-RateLimit-Reset"] = "60"
    response.headers["Retry-After"] = "60"

    return response


app.include_router(users_v1.router, prefix="/api/v1/users", tags=["users-v1"])
app.include_router(tasks_v1.router, prefix="/api/v1/tasks", tags=["tasks-v1"])

app.include_router(users_v2.router, prefix="/api/v2/users", tags=["users-v2"])
app.include_router(tasks_v2.router, prefix="/api/v2/tasks", tags=["tasks-v2"])

app.include_router(internal_tasks.router, prefix="/internal/tasks", tags=["internal_tasks"])
app.include_router(internal_stats.router, prefix="/internal/stats", tags=["internal_stats"])

@app.get("/")
async def root():
    return {
        "message": "Task Management API v1.0.0",
        "docs": "http://localhost:8000/docs",
        "redoc": "http://localhost:8000/redoc",
        "version": "1.0.0"
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)