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

2026-02-24 11:28:25 +01:00
parent 5669536c55
commit 5444ee97e3
5 changed files with 247 additions and 101 deletions
@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from aare.common.logger_config import setup_logger
+
+logger = setup_logger("aareDAQ")
+
+class TransformationInvalidException(Exception):
+    def __init__(self, message: str = "Transformation is not implemented"):
+        super().__init__(message)
+        self.message = message
+        logger.error(message, extra={"exception:": Exception})
+
+    def __str__(self) -> str:
+        return self.message
+
+
+class LoopCenteringFailed(Exception):
+    def __init__(self, message: str = "Loop Centering did not detect a sample"):
+        super().__init__(message)
+        self.message = message
+        logger.error(message, extra={"exception:": Exception})
+
+    def __str__(self) -> str:
+        return self.message
+
+
+class MountingFailed(Exception):
+    def __init__(self, message: str = "A sample was not mounted"):
+        super().__init__(message)
+        self.message = message
+        logger.error(message, extra={"exception:": Exception})
+
+    def __str__(self) -> str:
+        return self.message
+
+
+class WarningTellException(Exception):
+    def __init__(self, message: str = "Warning error in TELL"):
+        super().__init__(message)
+        self.message = message
+        logger.error(message, extra={"exception:": Exception})
+
+    def __str__(self) -> str:
+        return self.message
+
+
+class CriticalTellException(Exception):
+    def __init__(self, message: str = "Critical error in TELL"):
+        super().__init__(message)
+        self.message = message
+        logger.error(message, extra={"exception:": Exception})
+
+    def __str__(self) -> str:
+        return self.message
+
+
+class AXCFailed(Exception):
+    def __init__(self, message: str = "Auto X-ray centering failed"):
+        super().__init__(message)
+        self.message = message
+        logger.error(message, extra={"exception:": Exception})
+
+    def __str__(self) -> str:
+        return self.message
+
+class BeamlineBusyException(Exception):
+    def __init__(self, message: str = "Beamline is in busy state"):
+        super().__init__(message)
+        self.message = message
+        logger.error(message, extra={"exception:": Exception})
+
+    def __str__(self) -> str:
+        return self.message
+
+class SampleException(Exception):
+    def __init__(self, message: str = "Sample not found"):
+        super().__init__(message)
+        self.message = message
+        logger.error(message, extra={"exception:": Exception})
+
+    def __str__(self) -> str:
+        return self.message
+
+class AuthenticationException(Exception):
+    def __init__(self, message: str = "Sample belongs to a different user."):
+        super().__init__(message)
+        self.message = message
+        logger.error(message, extra={"exception:": Exception})
+
+    def __str__(self) -> str:
+        return self.message
@@ -21,6 +21,8 @@ from aare.common.models import (
 from aare.common.beamline import MXBeamline
 from aare.common.logger_config import setup_logger

+from aare.common.exception_handler import BeamlineBusyException
+
 ABR_POS_ALIGN_DEF = Coordinate(x=-18, y=-0.266, z=0)
 ABR_POS_MOUNT = Coordinate(x=18, y=0, z=0)#Coordinate(x=-18, y=0, z=0)
 ABR_OMEGA_MOUNT = 0.0
@@ -47,10 +49,6 @@ def base64_to_numpy(encoded_str: str | None) -> np.ndarray | None:
    return np.load(buffer)  # Load buffer as a NumPy array


-class BeamlineBusyException(Exception):
-    pass
-
-
 class BeamlineConfig:
    """
    Manages the configuration and state of a beamline system by interacting with a Redis
@@ -35,6 +35,15 @@ from aare.common.sample_geometry import SampleGeometryModel
 from aare.devices.area_detector import AutoEnum
 from aare.devices.jfjoch import JFJochWrapper

+from aare.common.exception_handler import (
+    TransformationInvalidException,
+    LoopCenteringFailed,
+    MountingFailed,
+    WarningTellException,
+    CriticalTellException,
+    AXCFailed,
+)
+
 logger = setup_logger("aareDAQ")

 class TransformationInvalidException(Exception):
@@ -24,10 +24,20 @@ from starlette.responses import StreamingResponse
 from aare.daq import auth
 from aare.common.beamline import mx_beamline
 from aare.daq.config import BeamlineConfig
-from aare.daq.daq import (AareDAQ, LoopCenteringFailed, TransformationInvalidException,
-                              MountingFailed, WarningTellException, CriticalTellException)
+from aare.daq.daq import AareDAQ
+
+from aare.daq.server_exception_handler import register_exception_handlers
+
+from aare.common.exception_handler import (
+    LoopCenteringFailed,
+    TransformationInvalidException,
+    MountingFailed,
+    WarningTellException,
+    CriticalTellException, SampleException, AuthenticationException,
+)

 app = FastAPI()
+register_exception_handlers(app)

 # OAuth2 setup
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
@@ -51,34 +61,25 @@ async def login(form_data: OAuth2PasswordRequestForm = Depends()):
 async def status(token: str = Depends(oauth2_scheme)) -> DAQStatusModel:
    data = auth.parse_token(token)

-    try:
-        full = daq.status
-        active_pgroup = cfg.pgroup
-        is_staff = data.staff
-        in_allowed_groups = (active_pgroup is not None and active_pgroup in data.pgroups)
-        in_ro = is_staff or in_allowed_groups
-        sample_pgroup = full.sample.user if full.sample is not None else None
-        sample_view_allowed = sample_pgroup in data.pgroups or is_staff
-
-        full.sample   = full.sample if in_ro and sample_view_allowed else None
-        full.box = full.box if in_ro else None
-        full.last_best_res = full.last_best_res if in_ro else None
-        full.last_best_b_factor = full.last_best_b_factor if in_ro else None
-        full.crystal_size = full.crystal_size if in_ro else CrystalSize(x=0,y=0,z=0)
-        full.session = SessionStatus(
-                current_pgroup=cfg.pgroup,
-                session=cfg.session_state(data.session),
-                staff=data.staff
-        )
-        return full
-
-    except Exception as e:
-        logger.error(f"Error getting status: {e}")
-        raise HTTPException(
-            status_code=api_status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"Error getting status: {e}"
-        )
+    full = daq.status
+    active_pgroup = cfg.pgroup
+    is_staff = data.staff
+    in_allowed_groups = (active_pgroup is not None and active_pgroup in data.pgroups)
+    in_ro = is_staff or in_allowed_groups
+    sample_pgroup = full.sample.user if full.sample is not None else None
+    sample_view_allowed = sample_pgroup in data.pgroups or is_staff

+    full.sample             = full.sample if in_ro and sample_view_allowed else None
+    full.box                = full.box if in_ro else None
+    full.last_best_res      = full.last_best_res if in_ro else None
+    full.last_best_b_factor = full.last_best_b_factor if in_ro else None
+    full.crystal_size       = full.crystal_size if in_ro else CrystalSize(x=0,y=0,z=0)
+    full.session = SessionStatus(
+            current_pgroup=cfg.pgroup,
+            session=cfg.session_state(data.session),
+            staff=data.staff
+    )
+    return full

@app.get("/beamline/geometry")
 async def sample_geometry(
@@ -262,41 +263,14 @@ async def mount(dbid: int, token: str = Depends(oauth2_scheme), reference: bool
            index = i

    if index == -1:
-        raise HTTPException(
-            status_code=api_status.HTTP_404_NOT_FOUND,
-            detail="Sample not found",
-        )
-    if token_data.staff or st.s[index].user in token_data.pgroups:
-        try:
-            daq.sample = st.s[index]
-        except MountingFailed as e:
-            raise HTTPException(
-                status_code=api_status.HTTP_404_NOT_FOUND,
-                detail=f"{e}",
-            )
-        except WarningTellException as e:
-            raise HTTPException(
-                status_code=api_status.HTTP_410_GONE,
-                detail=f"{e}",
-            )
-        except CriticalTellException as e:
-            raise HTTPException(
-                status_code=api_status.HTTP_417_EXPECTATION_FAILED,
-                detail=f"{e}"
-            )
-        except Exception as e:
-            raise HTTPException(
-                status_code=api_status.HTTP_500_INTERNAL_SERVER_ERROR,
-                detail=f"{e}"
-            )
-        return "OK"
-    else:
-        raise HTTPException(
-            status_code=api_status.HTTP_401_UNAUTHORIZED,
-            detail="Sample belongs to a different user.",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
+        raise SampleException(message="Sample not found")

+    if not (token_data.staff or st.s[index].user in token_data.pgroups):
+        raise AuthenticationException(message="Sample belongs to a different user.")
+
+    daq.sample = st.s[index]
+
+    return "OK"

@app.post("/sample/unmount")
 async def unmount(token: str = Depends(oauth2_scheme)):
@@ -418,40 +392,7 @@ async def rotation(val: RotationScanRequest, token: str = Depends(oauth2_scheme)
@app.post("/scan/auto")
 async def auto(s: SampleShortInfo, token: str = Depends(oauth2_scheme)):
    auth.check_jwt_rw(cfg, auth.parse_token(token))
-    try:
-        runtime = daq.measure(s)
-        return f"{runtime:0.3f}"
-    except LoopCenteringFailed as e:
-        raise HTTPException(
-            status_code=api_status.HTTP_404_NOT_FOUND,
-            detail=f"Loop centering failed: {e}",
-        )
-    except TransformationInvalidException as e:
-        raise HTTPException(
-            status_code=api_status.HTTP_400_BAD_REQUEST,
-            detail=f"Transformation invalid: {e}",
-        )
-    except MountingFailed as e:
-        raise HTTPException(
-            status_code=api_status.HTTP_404_NOT_FOUND,
-            detail=f"{e}",
-        )
-    except WarningTellException as e:
-        raise HTTPException(
-            status_code=api_status.HTTP_410_GONE,
-            detail=f"{e}",
-        )
-    except CriticalTellException as e:
-        raise HTTPException(
-            status_code=api_status.HTTP_417_EXPECTATION_FAILED,
-            detail=f"{e}"
-        )
-    except Exception as e:
-        logger.error(f"Exception in auto: {e}")
-        raise HTTPException(
-            status_code=api_status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"error {e}"
-        )
+    runtime = daq.measure(s)
    return f"{runtime:0.3f}"


@@ -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"),
+        )