Updated exception hadnling in server to use fastAPI exception handlers, and moved custom exception classes to exception_handler.py in Common

src/aare/daq/server_exception_handler.py

@@ -0,0 +1,107 @@
+# 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
+)
+
+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=api_status.HTTP_401_UNAUTHORIZED,
+            content=_error_payload(code="AUTHENTICATION_ERROR", message=str(exc) or "Invalid authentication"),
+        )
+
+    @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(Exception)
+    async def unhandled_exception_handler(request: Request, exc: Exception) -> JSONResponse:
+        logger.exception("Unhandled server exception")
+        return JSONResponse(
+            status_code=api_status.HTTP_500_INTERNAL_SERVER_ERROR,
+            content=_error_payload(code="INTERNAL_SERVER_ERROR", message="Internal server error"),
+        )