AareDAQ/src/aare/daq/server_exception_handler.py

# aare/common/server_error_handler.py
from __future__ import annotations

from fastapi import HTTPException
from fastapi import status as api_status
from starlette.requests import Request
from starlette.responses import JSONResponse

from aare.common.logger_config import setup_logger
from aare.common.exception_handler import (
    MountingFailed,
    WarningTellException,
    CriticalTellException,
    LoopCenteringFailed,
    TransformationInvalidException,
    BeamlineBusyException,
    AuthenticationException,
    SampleException,
    UserRightsException,
    SmargonCommunicationError,
    TellCommunicationError,
    JFJochCommunicationError,
    AerotechCommunicationError,
)

logger = setup_logger("aareDAQ")


def _error_payload(*, code: str, message: str, extra: dict | None = None) -> dict:
    payload = {"code": code, "message": message}
    if extra:
        payload["extra"] = extra
    return payload


def register_exception_handlers(app) -> None:
    """
    Register server-wide exception handlers on the given FastAPI app.
    Call once right after `app = FastAPI()`.
    """

    @app.exception_handler(HTTPException)
    async def http_exception_handler(request: Request, exc: HTTPException) -> JSONResponse:
        # Keep explicit HTTP errors, but normalize response shape
        detail = exc.detail
        if isinstance(detail, dict) and "code" in detail and "message" in detail:
            body = detail
        else:
            body = _error_payload(code="HTTP_ERROR", message=str(detail))
        return JSONResponse(status_code=exc.status_code, content=body, headers=exc.headers)

    @app.exception_handler(MountingFailed)
    async def mounting_failed_handler(request: Request, exc: MountingFailed) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_404_NOT_FOUND,
            content=_error_payload(code="MOUNTING_FAILED", message=str(exc)),
        )

    @app.exception_handler(WarningTellException)
    async def warning_tell_handler(request: Request, exc: WarningTellException) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_410_GONE,
            content=_error_payload(code="TELL_WARNING", message=str(exc)),
        )

    @app.exception_handler(CriticalTellException)
    async def critical_tell_handler(request: Request, exc: CriticalTellException) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_417_EXPECTATION_FAILED,
            content=_error_payload(code="TELL_CRITICAL", message=str(exc)),
        )

    @app.exception_handler(LoopCenteringFailed)
    async def loop_centering_failed_handler(request: Request, exc: LoopCenteringFailed) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_404_NOT_FOUND,
            content=_error_payload(code="LOOP_CENTERING_FAILED", message=str(exc)),
        )

    @app.exception_handler(TransformationInvalidException)
    async def transformation_invalid_handler(request: Request, exc: TransformationInvalidException) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_400_BAD_REQUEST,
            content=_error_payload(code="TRANSFORMATION_INVALID", message=str(exc)),
        )

    @app.exception_handler(BeamlineBusyException)
    async def beamline_busy_handler(request: Request, exc: BeamlineBusyException) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_423_LOCKED,
            content=_error_payload(code="BEAMLINE_BUSY", message=str(exc) or "Beamline is busy"),
        )

    @app.exception_handler(AuthenticationException)
    async def authentication_exception_handler(request: Request, exc: AuthenticationException) -> JSONResponse:
        return JSONResponse(
            status_code=getattr(exc, "status_code", api_status.HTTP_401_UNAUTHORIZED),
            content=_error_payload(
                code=str(getattr(exc, "code", "AUTHENTICATION_ERROR")),
                message=str(exc) or "Invalid authentication",
            ),
            headers=getattr(exc, "headers", None),
        )

    @app.exception_handler(UserRightsException)
    async def user_rights_exception_handler(request: Request, exc: UserRightsException) -> JSONResponse:
        return JSONResponse(
            status_code=getattr(exc, "status_code", api_status.HTTP_403_FORBIDDEN),
            content=_error_payload(
                code=str(getattr(exc, "code", "FORBIDDEN")),
                message=str(exc) or "Forbidden",
            ),
            headers=getattr(exc, "headers", None),
        )

    @app.exception_handler(SampleException)
    async def sample_exception_handler(request: Request, exc: SampleException) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_404_NOT_FOUND,
            content=_error_payload(code="SAMPLE_NOT_FOUND", message=str(exc) or "Sample not found"),
        )

    @app.exception_handler(SmargonCommunicationError)
    async def smargon_comm_handler(request: Request, exc: SmargonCommunicationError) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_503_SERVICE_UNAVAILABLE,
            content=_error_payload(
                code="SMARGON_UNAVAILABLE",
                message=str(exc) or "Smargon is unavailable",
                extra={
                    "operation": getattr(exc, "operation", None),
                    "endpoint": getattr(exc, "endpoint", None),
                },
            ),
        )

    @app.exception_handler(TellCommunicationError)
    async def tell_comm_handler(request: Request, exc: TellCommunicationError) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_503_SERVICE_UNAVAILABLE,
            content=_error_payload(
                code="TELL_UNAVAILABLE",
                message=str(exc) or "TELL is unavailable",
                extra={
                    "operation": getattr(exc, "operation", None),
                    "endpoint": getattr(exc, "endpoint", None),
                },
            ),
        )

    @app.exception_handler(JFJochCommunicationError)
    async def jfjoch_comm_handler(request: Request, exc: JFJochCommunicationError) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_503_SERVICE_UNAVAILABLE,
            content=_error_payload(
                code="JFJOCH_UNAVAILABLE",
                message=str(exc) or "JFJoch detector is unavailable",
                extra={
                    "operation": getattr(exc, "operation", None),
                    "endpoint": getattr(exc, "endpoint", None),
                    "base_url": getattr(exc, "base_url", None),
                },
            ),
        )

    @app.exception_handler(AerotechCommunicationError)
    async def aerotech_comm_handler(request: Request, exc: AerotechCommunicationError) -> JSONResponse:
        return JSONResponse(
            status_code=api_status.HTTP_503_SERVICE_UNAVAILABLE,
            content=_error_payload(
                code="AEROTECH_UNAVAILABLE",
                message=str(exc) or "Aerotech is unavailable",
                extra={
                    "operation": getattr(exc, "operation", None),
                    "endpoint": getattr(exc, "endpoint", None),
                    "base_url": getattr(exc, "base_url", None),
                },
            ),
        )

    @app.exception_handler(Exception)
    async def unhandled_exception_handler(request: Request, exc: Exception) -> JSONResponse:
        logger.exception(f"Unhandled server exception: {exc}")
        return JSONResponse(
            status_code=api_status.HTTP_500_INTERNAL_SERVER_ERROR,
            content=_error_payload(code="INTERNAL_SERVER_ERROR", message=str(exc) or "Internal server error"),
        )