be-optimumoh/src/exceptions.py

# Define base error model
import logging
from typing import Any, Dict, List, Optional

from asyncpg.exceptions import DataError as AsyncPGDataError
from asyncpg.exceptions import PostgresError
from fastapi import FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from slowapi import _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from sqlalchemy.exc import (DataError, DBAPIError, IntegrityError,
                            SQLAlchemyError)

from src.enums import ResponseStatus


class ErrorDetail(BaseModel):
    field: Optional[str] = None
    message: str
    code: Optional[str] = None
    params: Optional[Dict[str, Any]] = None


class ErrorResponse(BaseModel):
    data: Optional[Any] = None
    message: str
    status: ResponseStatus = ResponseStatus.ERROR
    errors: Optional[List[ErrorDetail]] = None


# Custom exception handler setup


def get_request_context(request: Request):
    """
    Get detailed request context for logging.
    """

    def get_client_ip():
        """
        Get the real client IP address from Kong Gateway headers.
        Kong sets X-Real-IP and X-Forwarded-For headers by default.
        """
        # Kong specific headers
        if "X-Real-IP" in request.headers:
            return request.headers["X-Real-IP"]

        # Fallback to X-Forwarded-For
        if "X-Forwarded-For" in request.headers:
            # Get the first IP (original client)
            return request.headers["X-Forwarded-For"].split(",")[0].strip()

        # Last resort
        return request.client.host

    return {
        "endpoint": request.url.path,
        "url": request.url,
        "method": request.method,
        "remote_addr": get_client_ip(),
    }


def handle_sqlalchemy_error(error: SQLAlchemyError):
    """
    Handle SQLAlchemy errors and return user-friendly error messages.
    """
    original_error = getattr(error, "orig", None)
    print(original_error)

    if isinstance(error, IntegrityError):
        if "unique constraint" in str(error).lower():
            return "This record already exists.", 409
        elif "foreign key constraint" in str(error).lower():
            return "Related record not found.", 400
        else:
            return "Data integrity error.", 400
    elif isinstance(error, DataError) or isinstance(original_error, AsyncPGDataError):
        return "Invalid data provided.", 400
    elif isinstance(error, DBAPIError):
        if "unique constraint" in str(error).lower():
            return "This record already exists.", 409
        elif "foreign key constraint" in str(error).lower():
            return "Related record not found.", 400
        elif "null value in column" in str(error).lower():
            return "Required data missing.", 400
        elif "invalid input for query argument" in str(error).lower():
            return "Invalid data provided.", 400
        else:
            return "Database error.", 500
    else:
        # Log the full error for debugging purposes
        logging.error(f"Unexpected database error: {str(error)}")
        return "An unexpected database error occurred.", 500


def handle_exception(request: Request, exc: Exception):
    """
    Global exception handler for Fastapi application.
    """
    request_info = get_request_context(request)

    if isinstance(exc, RateLimitExceeded):
        _rate_limit_exceeded_handler(request, exc)
    if isinstance(exc, HTTPException):
        logging.error(
            f"HTTP exception | Code: {exc.status_code} | Error: {exc.detail} | Request: {request_info}",
            extra={"error_category": "http"},
        )

        return JSONResponse(
            status_code=exc.status_code,
            content={
                "data": None,
                "message": str(exc.detail),
                "status": ResponseStatus.ERROR,
                "errors": [ErrorDetail(message=str(exc.detail)).model_dump()],
            },
        )

    if isinstance(exc, SQLAlchemyError):
        error_message, status_code = handle_sqlalchemy_error(exc)
        logging.error(
            f"Database Error | Error: {str(error_message)} | Request: {request_info}",
            extra={"error_category": "database"},
        )

        return JSONResponse(
            status_code=status_code,
            content={
                "data": None,
                "message": error_message,
                "status": ResponseStatus.ERROR,
                "errors": [ErrorDetail(message=error_message).model_dump()],
            },
        )

    # Log unexpected errors
    logging.error(
        f"Unexpected Error | Error: {str(exc)} | Request: {request_info}",
        extra={"error_category": "unexpected"},
    )

    return JSONResponse(
        status_code=500,
        content={
            "data": None,
            "message": str(exc),
            "status": ResponseStatus.ERROR,
            "errors": [
                ErrorDetail(message="An unexpected error occurred.").model_dump()
            ],
        },
    )