From 5c433079332edbab7cb1b9d1b075101bf3214391 Mon Sep 17 00:00:00 2001 From: adam Date: Mon, 28 Oct 2024 21:31:45 -0700 Subject: [PATCH] finish business --- src/mobu/exceptions.py | 54 ++++++++++++++++++++++----- src/mobu/services/business/nublado.py | 35 +++++------------ 2 files changed, 54 insertions(+), 35 deletions(-) diff --git a/src/mobu/exceptions.py b/src/mobu/exceptions.py index 29969af3..756ec50a 100644 --- a/src/mobu/exceptions.py +++ b/src/mobu/exceptions.py @@ -70,11 +70,6 @@ def _remove_ansi_escapes(string: str) -> str: class MobuMixin: """Mixin class to add `event` and `monkey` fields to Exception.""" - def __init__( - self, event: str | None = None, monkey: str | None = None - ) -> None: - self.mobu_init() - def mobu_init( self, event: str | None = None, monkey: str | None = None ) -> None: @@ -221,13 +216,52 @@ def __init__( self.mobu_init(monkey=monkey, event=event) @classmethod - def from_slack_exception(cls, exc: ne.NubladoClientSlackException) -> Self: - return cls( + def from_client_exception( + cls, + exc: ne.NubladoClientSlackException, + *, + started_at: datetime.datetime | None = None, + failed_at: datetime.datetime | None = None, + monkey: str | None = None, + event: str | None = None, + annotations: dict[str, str] | None = None, + ) -> Self: + """ + Add Mobu-specific fields to exception from NubladoClient layer. + + Parameters + ---------- + exc + Original exception + started_at + Timestamp for beginning of operation that caused the exception, + if known. + failed_at + Timestamp for failure of operation that caused the exception, + if known (defaults to the current time). + monkey + Monkey spawning the lab, if known. + event + Event (from mobu's perspective) spawning the lab, if known. + annotations + Additional annotations + + Returns + ------- + MobuSlackException + Converted exception. + """ + new_exc = cls( msg=exc.message, user=exc.user, - started_at=exc.started_at, - failed_at=exc.failed_at, + started_at=started_at or exc.started_at, + failed_at=failed_at or exc.failed_at, + event=event, + monkey=monkey, ) + new_exc.annotations.update(exc.annotations or {}) + new_exc.annotations.update(annotations or {}) + return new_exc @override def common_fields(self) -> list[SlackBaseField]: @@ -267,7 +301,7 @@ def to_slack(self) -> SlackMessage: class MobuSlackWebException( ne.NubladoClientSlackWebException, MobuSlackException ): - """Represents an exception that can be reported to Slack. + """Represents a web exception that can be reported to Slack. Similar to `MobuSlackException`, this adds some additional fields to `~rubin.nublado.client.SlackWebException` but is otherwise equivalent. It diff --git a/src/mobu/services/business/nublado.py b/src/mobu/services/business/nublado.py index ad053cb8..bb2980f7 100644 --- a/src/mobu/services/business/nublado.py +++ b/src/mobu/services/business/nublado.py @@ -10,17 +10,9 @@ from random import SystemRandom from typing import Generic, TypeVar +import rubin.nublado.client.exceptions as ne from httpx import AsyncClient from rubin.nublado.client import JupyterLabSession, NubladoClient -from rubin.nublado.client.exceptions import ( - CodeExecutionError as ClientCodeExecutionError, -) -from rubin.nublado.client.exceptions import ( - JupyterProtocolError as ClientJupyterProtocolError, -) -from rubin.nublado.client.exceptions import ( - JupyterWebError as ClientJupyterWebError, -) from safir.datetime import current_datetime, format_datetime_for_logging from safir.slack.blockkit import SlackException from structlog.stdlib import BoundLogger @@ -32,7 +24,6 @@ JupyterSpawnError, JupyterTimeoutError, JupyterWebError, - MobuMixin, ) from ...models.business.nublado import ( NubladoBusinessData, @@ -188,20 +179,14 @@ async def execute(self) -> None: try: await self._execute() except Exception as exc: - monkey = None - event = "execute_code" # Fallback - if isinstance(exc, MobuMixin): - if exc.monkey: - monkey = exc.monkey - if exc.event: - event = exc.event - if isinstance(exc, ClientCodeExecutionError): + monkey = getattr(exc, "monkey", None) + event = getattr(exc, "event", "execute_code") + if isinstance(exc, ne.CodeExecutionError): raise CodeExecutionError.from_client_exception( exc, monkey=monkey, event=event, annotations=self.annotations(), - started_at=exc.started_at, ) from exc raise @@ -245,14 +230,14 @@ async def hub_login(self) -> None: with self.timings.start("hub_login", self.annotations()) as sw: try: await self._client.auth_to_hub() - except ClientJupyterProtocolError as exc: + except ne.JupyterProtocolError as exc: raise JupyterProtocolError.from_client_exception( exc, event=sw.event, annotations=sw.annotations, started_at=sw.start_time, ) from exc - except ClientJupyterWebError as exc: + except ne.JupyterWebError as exc: raise JupyterWebError.from_client_exception( exc, event=sw.event, @@ -265,7 +250,7 @@ async def spawn_lab(self) -> bool: # noqa: C901 timeout = self.options.spawn_timeout try: await self._client.spawn_lab(self.options.image) - except ClientJupyterWebError as exc: + except ne.JupyterWebError as exc: raise JupyterWebError.from_client_exception( exc, event=sw.event, @@ -296,7 +281,7 @@ async def spawn_lab(self) -> bool: # noqa: C901 event=sw.event, started_at=sw.start_time, ) from None - except ClientJupyterWebError as exc: + except ne.JupyterWebError as exc: raise JupyterWebError.from_client_exception( exc, event=sw.event, @@ -344,7 +329,7 @@ async def lab_login(self) -> None: with self.timings.start("lab_login", self.annotations()) as sw: try: await self._client.auth_to_lab() - except ClientJupyterProtocolError as exc: + except ne.JupyterProtocolError as exc: raise JupyterProtocolError.from_client_exception( exc, event=sw.event, @@ -401,7 +386,7 @@ async def delete_lab(self) -> None: with self.timings.start("delete_lab", self.annotations()) as sw: try: await self._client.stop_lab() - except ClientJupyterWebError as exc: + except ne.JupyterWebError as exc: raise JupyterWebError.from_client_exception( exc, event=sw.event,