api.py

from uuid import UUID
import base64
from fastapi import FastAPI, Request, Depends, Header
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import Annotated, Optional

from app.database.database import Database
from app.repository import Repository
from app.service import Service
from app.exceptions import (
    MissingAuthorizationHeaderException,
    InvalidCredentialsException,
    InvalidAuthorizationHeaderFormatException,
    UserNotFoundException,
    ProductNotFoundException,
    InsufficientProductQuantityException,
    InsufficientBalanceException,
)
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry import trace
from opentelemetry.trace import StatusCode
from app.tracing import configure_provider
from app.log import configure_logging, get_logger
import gc

configure_provider("python-tracing-example")
configure_logging()

app = FastAPI()
database = Database()
repository = Repository(database)
service = Service(repository)

FastAPIInstrumentor.instrument_app(app)
tracer = trace.get_tracer(__file__)
logger = get_logger(__name__)


@app.exception_handler(Exception)
async def generic_exception_handler(request: Request, exc: Exception):
    span = trace.get_current_span()
    span.record_exception(exc)
    span.set_status(StatusCode.ERROR)
    logger.error(
        "Internal server error",
        exc_info=exc,
    )
    return JSONResponse(
        status_code=500,
        content={"detail": "Internal server error"},
    )


@app.exception_handler(MissingAuthorizationHeaderException)
async def missing_authorization_header_handler(
    request: Request, exc: MissingAuthorizationHeaderException
):
    span = trace.get_current_span()
    span.record_exception(exc)
    span.set_status(StatusCode.ERROR)
    return JSONResponse(
        status_code=401,
        content={"detail": str(exc)},
        headers={"WWW-Authenticate": "Basic"},
    )


@app.exception_handler(InvalidCredentialsException)
async def invalid_credentials_handler(
    request: Request, exc: InvalidCredentialsException
):
    span = trace.get_current_span()
    span.record_exception(exc)
    span.set_status(StatusCode.ERROR)
    return JSONResponse(
        status_code=401,
        content={"detail": str(exc)},
        headers={"WWW-Authenticate": "Basic"},
    )


@app.exception_handler(InvalidAuthorizationHeaderFormatException)
async def invalid_authorization_header_format_handler(
    request: Request, exc: InvalidAuthorizationHeaderFormatException
):
    span = trace.get_current_span()
    span.record_exception(exc)
    span.set_status(StatusCode.ERROR)
    return JSONResponse(
        status_code=401,
        content={"detail": str(exc)},
        headers={"WWW-Authenticate": "Basic"},
    )


@app.exception_handler(UserNotFoundException)
async def user_not_found_handler(request: Request, exc: UserNotFoundException):
    span = trace.get_current_span()
    span.record_exception(exc)
    span.set_status(StatusCode.ERROR)
    return JSONResponse(
        status_code=404,
        content={"detail": str(exc)},
    )


@app.exception_handler(ProductNotFoundException)
async def product_not_found_handler(request: Request, exc: ProductNotFoundException):
    span = trace.get_current_span()
    span.record_exception(exc)
    span.set_status(StatusCode.ERROR)
    return JSONResponse(
        status_code=404,
        content={"detail": str(exc)},
    )


@app.exception_handler(InsufficientProductQuantityException)
async def insufficient_product_quantity_handler(
    request: Request, exc: InsufficientProductQuantityException
):
    span = trace.get_current_span()
    span.record_exception(exc)
    span.set_status(StatusCode.ERROR)
    return JSONResponse(
        status_code=400,
        content={"detail": str(exc)},
    )


@app.exception_handler(InsufficientBalanceException)
async def insufficient_balance_handler(
    request: Request, exc: InsufficientBalanceException
):
    span = trace.get_current_span()
    span.record_exception(exc)
    span.set_status(StatusCode.ERROR)
    return JSONResponse(
        status_code=400,
        content={"detail": str(exc)},
    )


class BuyProductPayload(BaseModel):
    quantity: int


class BuyProductResponse(BaseModel):
    message: str


@tracer.start_as_current_span("authentication")
async def authentication(
    authorization: Annotated[Optional[str], Header()] = None,
) -> UUID:
    """Dependency that validates Basic Auth credentials and returns the user UUID."""
    span = trace.get_current_span()
    if not authorization or not authorization.startswith("Basic "):
        span.add_event("auth header missing")
        span.set_status(StatusCode.ERROR)
        raise MissingAuthorizationHeaderException()

    try:
        # Extract and decode the base64 encoded credentials
        encoded_credentials = authorization.split(" ", 1)[1]
        decoded_credentials = base64.b64decode(encoded_credentials).decode("utf-8")

        # Parse user_uuid:password format
        user_uuid_str, password = decoded_credentials.split(":", 1)
        user_uuid = UUID(user_uuid_str)
        span.set_attribute("user.id", str(user_uuid))

        # Verify credentials using the service
        is_valid = await service.verify_user_password(user_uuid, password)

        if not is_valid:
            span.add_event("invalid user password")
            span.set_status(StatusCode.ERROR)
            raise InvalidCredentialsException()
        span.add_event("user successfully logged")

        return user_uuid

    except (ValueError, UnicodeDecodeError, IndexError):
        span.add_event("invalid authorization header format")
        span.set_status(StatusCode.ERROR)
        raise InvalidAuthorizationHeaderFormatException()


@app.middleware("http")
async def gc_stats(request: Request, call_next):
    with tracer.start_as_current_span("gc_stats"):
        span = trace.get_current_span()
        response = await call_next(request)
        stats = gc.get_stats()
        span.set_attribute("python.gc.0.collections", stats[0]["collections"])
        span.set_attribute("python.gc.1.collections", stats[1]["collections"])
        span.set_attribute("python.gc.2.collections", stats[2]["collections"])
        span.set_attribute("python.gc.0.collected", stats[0]["collected"])
        span.set_attribute("python.gc.1.collected", stats[1]["collected"])
        span.set_attribute("python.gc.2.collected", stats[2]["collected"])
        span.set_attribute("python.gc.0.uncollectable", stats[0]["uncollectable"])
        span.set_attribute("python.gc.1.uncollectable", stats[1]["uncollectable"])
        span.set_attribute("python.gc.2.uncollectable", stats[2]["uncollectable"])
        return response


@app.post("/product/{id}", status_code=200)
async def buy_product(
    id: UUID,
    payload: BuyProductPayload,
    user_id: Annotated[UUID, Depends(authentication)],
) -> BuyProductResponse:
    span = trace.get_current_span()
    span.set_attribute("user.id", str(user_id))
    span.set_attribute("product.id", str(id))
    await service.buy_product(product_id=id, user_id=user_id, quantity=payload.quantity)
    return BuyProductResponse(message="success")