diff --git a/src/aare/common/exception_handler.py b/src/aare/common/exception_handler.py new file mode 100644 index 00000000..35b6ef65 --- /dev/null +++ b/src/aare/common/exception_handler.py @@ -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 diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index 2873747b..c8b5abad 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -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 diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index b3feea28..8a806337 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -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): diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index f482271b..0ee2cffb 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -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}" diff --git a/src/aare/daq/server_exception_handler.py b/src/aare/daq/server_exception_handler.py new file mode 100644 index 00000000..da273fd6 --- /dev/null +++ b/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"), + )