feat: updated error handling in job status retrieval

JanssenBrm · JanssenBrm · commit 9ac6d09b40b3 · 2025-12-10T14:00:28.000+01:00
diff --git a/app/auth.py b/app/auth.py
@@ -6,7 +6,8 @@
 from jwt import PyJWKClient
 from loguru import logger
 
-from app.error import AuthException
+from app.error import AuthException, DispatcherException
+from app.schemas.websockets import WSStatusMessage
 
 from .config.settings import settings
 
@@ -57,6 +58,7 @@ async def websocket_authenticate(websocket: WebSocket) -> str | None:
     """
     logger.debug("Authenticating websocket")
     token = websocket.query_params.get("token")
+
     if not token:
         logger.error("Token is missing from websocket authentication")
         await websocket.close(code=1008, reason="Missing token")
@@ -65,9 +67,20 @@ async def websocket_authenticate(websocket: WebSocket) -> str | None:
     try:
         await websocket.accept()
         return token
+    except DispatcherException as ae:
+        logger.error(f"Dispatcher exception detected: {ae.message}")
+        await websocket.send_json(
+            WSStatusMessage(type="error", message=ae.message).model_dump()
+        )
+        await websocket.close(code=1008, reason=ae.error_code)
+        return None
     except Exception as e:
-        logger.error(f"Invalid token in websocket authentication: {e}")
-        await websocket.close(code=1008, reason="Invalid token")
+        logger.error(f"Unexpected error occurred during websocket authentication: {e}")
+        await WSStatusMessage(
+            type="error",
+            message="Something went wrong during authentication. Please try again.",
+        ).model_dump()
+        await websocket.close(code=1008, reason="INTERNAL_ERROR")
         return None
 
 
@@ -107,9 +120,12 @@ async def exchange_token_for_provider(
             resp = await client.post(token_url, data=payload)
     except httpx.RequestError as exc:
         logger.error(f"Token exchange network error for provider={provider}: {exc}")
-        raise HTTPException(
-            status_code=status.HTTP_502_BAD_GATEWAY,
-            detail="Failed to contact the identity provider for token exchange.",
+        raise AuthException(
+            http_status=status.HTTP_502_BAD_GATEWAY,
+            message=(
+                f"Could not authenticate with {provider}. Please contact APEx support or reach out "
+                "through the <a href='https://forum.apex.esa.int/'>APEx User Forum</a>."
+            ),
         )
 
     # Parse response
@@ -119,9 +135,12 @@ async def exchange_token_for_provider(
         logger.error(
             f"Token exchange invalid JSON response (status={resp.status_code})"
         )
-        raise HTTPException(
-            status_code=status.HTTP_502_BAD_GATEWAY,
-            detail="Invalid response from identity provider during token exchange.",
+        raise AuthException(
+            http_status=status.HTTP_502_BAD_GATEWAY,
+            message=(
+                f"Could not authenticate with {provider}. Please contact APEx support or reach out "
+                "through the <a href='https://forum.apex.esa.int/'>APEx User Forum</a>."
+            ),
         )
 
     if resp.status_code != 200:
@@ -138,7 +157,14 @@ async def exchange_token_for_provider(
             else status.HTTP_502_BAD_GATEWAY
         )
 
-        raise HTTPException(client_status, detail=body)
+        raise AuthException(
+            http_status=client_status,
+            message=(
+                f"Please link your account with {provider} in your <a href='https://{settings.keycloak_host}/realms/{settings.keycloak_realm}/account'>Account Dashboard</a>"
+                if body.get("error", "") == "not_linked"
+                else f"Could not authenticate with {provider}: {err}"
+            ),
+        )
 
     # Successful exchange, return token response (access_token, expires_in, etc.)
     return body
diff --git a/app/middleware/error_handling.py b/app/middleware/error_handling.py
@@ -1,4 +1,5 @@
 from fastapi import Request, status
+from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
 from app.error import DispatcherException, ErrorResponse
 from app.middleware.correlation_id import correlation_id_ctx
@@ -33,19 +34,31 @@ async def generic_exception_handler(request: Request, exc: Exception):
         request_id=correlation_id_ctx.get(),
     )
 
-    # Log exception to server logs for debugging
-    print(f"[ERROR] Request ID: {exc}")
-
+    logger.exception(f"GenericException raised: {exc}")
     return JSONResponse(
         status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=content.dict()
     )
 
 
+async def validation_exception_handler(request: Request, exc: RequestValidationError):
+
+    content = ErrorResponse(
+        error_code="VALIDATION_ERROR",
+        message="Request validation failed.",
+        details={"errors": exc.errors()},
+        request_id=correlation_id_ctx.get(),
+    )
+
+    return JSONResponse(
+        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, content=content.dict()
+    )
+
+
 def register_exception_handlers(app):
     """
     Call this in main.py after creating the FastAPI() instance.
     """
 
     app.add_exception_handler(DispatcherException, dispatch_exception_handler)
-    # app.add_exception_handler(RequestValidationError, validation_exception_handler)
+    app.add_exception_handler(RequestValidationError, validation_exception_handler)
     app.add_exception_handler(Exception, generic_exception_handler)
diff --git a/app/routers/jobs_status.py b/app/routers/jobs_status.py
@@ -7,6 +7,8 @@
 from loguru import logger
 
 from app.database.db import SessionLocal, get_db
+from app.error import DispatcherException, ErrorResponse, InternalException
+from app.middleware.error_handling import get_dispatcher_error_response
 from app.schemas.jobs_status import JobsFilter, JobsStatusResponse
 from app.schemas.websockets import WSStatusMessage
 from app.services.processing import get_processing_jobs_by_user_id
@@ -22,6 +24,19 @@
     "/jobs_status",
     tags=["Upscale Tasks", "Unit Jobs"],
     summary="Get a list of all upscaling tasks & processing jobs for the authenticated user",
+    responses={
+        InternalException.http_status: {
+            "description": "Internal server error",
+            "model": ErrorResponse,
+            "content": {
+                "application/json": {
+                    "example": get_dispatcher_error_response(
+                        InternalException(), "request-id"
+                    )
+                }
+            },
+        },
+    },
 )
 async def get_jobs_status(
     db: Session = Depends(get_db),
@@ -34,21 +49,29 @@ async def get_jobs_status(
     """
     Return combined list of upscaling tasks and processing jobs for the authenticated user.
     """
-    logger.debug("Fetching jobs list")
-    upscaling_tasks = (
-        await get_upscaling_tasks_by_user_id(token, db)
-        if JobsFilter.upscaling in filter
-        else []
-    )
-    processing_jobs = (
-        await get_processing_jobs_by_user_id(token, db)
-        if JobsFilter.processing in filter
-        else []
-    )
-    return JobsStatusResponse(
-        upscaling_tasks=upscaling_tasks,
-        processing_jobs=processing_jobs,
-    )
+    try:
+        logger.debug("Fetching jobs list")
+        upscaling_tasks = (
+            await get_upscaling_tasks_by_user_id(token, db)
+            if JobsFilter.upscaling in filter
+            else []
+        )
+        processing_jobs = (
+            await get_processing_jobs_by_user_id(token, db)
+            if JobsFilter.processing in filter
+            else []
+        )
+        return JobsStatusResponse(
+            upscaling_tasks=upscaling_tasks,
+            processing_jobs=processing_jobs,
+        )
+    except DispatcherException as de:
+        raise de
+    except Exception as e:
+        logger.error(f"Error retrieving job status: {e}")
+        raise InternalException(
+            message="An error occurred while retrieving the job status."
+        )
 
 
 @router.websocket(
@@ -91,8 +114,18 @@ async def ws_jobs_status(
 
     except WebSocketDisconnect:
         logger.info("WebSocket disconnected")
+    except DispatcherException as ae:
+        logger.error(f"Dispatcher exception detected: {ae.message}")
+        await websocket.send_json(
+            WSStatusMessage(type="error", message=ae.message).model_dump()
+        )
+        await websocket.close(code=1008, reason=ae.error_code)
     except Exception as e:
-        logger.exception(f"Error in jobs_status_ws: {e}")
-        await websocket.close(code=1011, reason="Error in job status websocket: {e}")
+        logger.error(f"Unexpected error occurred during websocket authentication: {e}")
+        await WSStatusMessage(
+            type="error",
+            message="Something went wrong during authentication. Please try again.",
+        ).model_dump()
+        await websocket.close(code=1008, reason="INTERNAL_ERROR")
     finally:
         db.close()
