From 2019f23eca041a7b6c6229ffec59e7ed60a6ef45 Mon Sep 17 00:00:00 2001
From: Jonathan Sick <jsick@lsst.org>
Date: Tue, 10 Sep 2024 13:53:06 -0400
Subject: [PATCH] Report exceptions in notebook execution in "error"

This modifies the result payload for a notebook execution job when the
job is complete.

- Add a new "error" field that reports an exception raised during the
  noteburst job, which explains why a job's "success" field is false.
  The field contains a NoteburstExecutionError model, which in turn
  contains an enum code for the error and a text message describing the
  error.

- Keep the existing "ipynb_error" field for reporting exceptions raised
  within the notebook itself. Noteburst's execution could still be
  considered successful even if the notebook itself raised an exception.
---
 src/noteburst/handlers/v1/models.py | 83 +++++++++++++++++++++++++++--
 1 file changed, 79 insertions(+), 4 deletions(-)

diff --git a/src/noteburst/handlers/v1/models.py b/src/noteburst/handlers/v1/models.py
index 3e61898..125e3dd 100644
--- a/src/noteburst/handlers/v1/models.py
+++ b/src/noteburst/handlers/v1/models.py
@@ -4,6 +4,7 @@
 
 import json
 from datetime import datetime, timedelta
+from enum import Enum
 from typing import Annotated, Any
 
 from arq.jobs import JobStatus
@@ -12,6 +13,7 @@
 from safir.arq import JobMetadata, JobResult
 from safir.pydantic import HumanTimedelta
 
+from noteburst.exceptions import NbexecTaskError, NbexecTaskTimeoutError
 from noteburst.jupyterclient.jupyterlab import (
     NotebookExecutionErrorModel,
     NotebookExecutionResult,
@@ -48,6 +50,34 @@ def from_nbexec_error(
         )
 
 
+class NoteburstErrorCodes(Enum):
+    """Error codes for Noteburst errors."""
+
+    timeout = "timeout"
+    """The notebook execution timed out."""
+
+    jupyter_error = "jupyter_error"
+    """An error occurred contacting the Jupyter server."""
+
+    unknown = "unknown"
+    """An unknown error occurred."""
+
+
+class NoteburstExecutionError(BaseModel):
+    """Information about an exception that occurred during noteburst's
+    execution of a notebook (other than an exception raised in the notebook
+    itself).
+    """
+
+    code: NoteburstErrorCodes = Field(
+        description="The reference code of the error."
+    )
+
+    message: str | None = Field(
+        None, description="Additional information about the exception."
+    )
+
+
 class NotebookResponse(BaseModel):
     """Information about a notebook execution job, possibly including the
     result and source notebooks.
@@ -102,6 +132,17 @@ class NotebookResponse(BaseModel):
         ),
     ] = None
 
+    error: Annotated[
+        NoteburstExecutionError | None,
+        Field(
+            description=(
+                "An error occurred during notebook execution, other than an "
+                "exception in the notebook itself. This field is null if an "
+                "error did not occur."
+            )
+        ),
+    ] = None
+
     ipynb: Annotated[
         str | None,
         Field(
@@ -130,18 +171,51 @@ async def from_job_metadata(
         job_result: JobResult | None = None,
     ) -> NotebookResponse:
         """Create a NotebookResponse from a job."""
+        # When a job is a "success" it means that the arq worker didn't raise
+        # an exception, so we can expect an ipynb result. However the ipynb
+        # might have still raised an exception which is part of
+        # nbexec_result.error and we want to pass that back to the user.
         if job_result is not None and job_result.success:
             nbexec_result = NotebookExecutionResult.model_validate_json(
                 job_result.result
             )
             ipynb = nbexec_result.notebook
             if nbexec_result.error:
-                error = NotebookError.from_nbexec_error(nbexec_result.error)
+                ipynb_error = NotebookError.from_nbexec_error(
+                    nbexec_result.error
+                )
             else:
-                error = None
+                ipynb_error = None
         else:
             ipynb = None
-            error = None
+            ipynb_error = None
+
+        # In this case the job is complete but failed (an exception was raised)
+        # so we want to pass the exception back to the user.
+        if job_result and not job_result.success:
+            try:
+                # Attempt to access the result to raise any exceptions
+                job_result.result  # noqa: B018
+            except NbexecTaskTimeoutError as e:
+                # format this error
+                noteburst_error = NoteburstExecutionError(
+                    code=NoteburstErrorCodes.timeout,
+                    message=str(e),
+                )
+            except NbexecTaskError as e:
+                # format this error
+                noteburst_error = NoteburstExecutionError(
+                    code=NoteburstErrorCodes.jupyter_error,
+                    message=str(e),
+                )
+            except Exception as e:
+                # format an unknown error
+                noteburst_error = NoteburstExecutionError(
+                    code=NoteburstErrorCodes.unknown,
+                    message=str(e),
+                )
+        else:
+            noteburst_error = None
 
         return cls(
             job_id=job.id,
@@ -153,8 +227,9 @@ async def from_job_metadata(
             start_time=job_result.start_time if job_result else None,
             finish_time=job_result.finish_time if job_result else None,
             success=job_result.success if job_result else None,
+            error=noteburst_error,
             ipynb=ipynb,
-            ipynb_error=error,
+            ipynb_error=ipynb_error,
         )