Raise SessionStatusError on get(_or_create) endpoint

Instead of manually handling specific state cases

Author: FlorentinD

changelog.md

@@ -16,6 +16,7 @@
 * Improve error message when a query is interrupted by a signal (SIGINT or SIGTERM).
 * Improve error message if session is expired.
 * Improve robustness of Arrow client against connection errors such as `FlightUnavailableError` and `FlightTimedOutError`.
+* Return dedicated error class `SessionStatusError` if a session failed or expired.
 
 
 ## Other changes

graphdatascience/session/__init__.py

@@ -1,4 +1,5 @@
 from .algorithm_category import AlgorithmCategory
+from .aura_api import AuraApiError, SessionStatusError
 from .cloud_location import CloudLocation
 from .dbms_connection_info import DbmsConnectionInfo
 from .gds_sessions import AuraAPICredentials, GdsSessions
@@ -14,4 +15,6 @@
     "SessionMemory",
     "SessionMemoryValue",
     "AlgorithmCategory",
+    "SessionStatusError",
+    "AuraApiError",
 ]

graphdatascience/session/aura_api.py

@@ -23,6 +23,8 @@
     InstanceDetails,
     InstanceSpecificDetails,
     SessionDetails,
+    SessionDetailsWithErrors,
+    SessionErrorData,
     TenantDetails,
     WaitResult,
 )
@@ -32,12 +34,27 @@
 
 
 class AuraApiError(Exception):
+    """
+    Raised when an API call to the AuraAPI fails (after retries).
+    """
+
     def __init__(self, message: str, status_code: int):
         super().__init__(self, message)
         self.status_code = status_code
         self.message = message
 
 
+class SessionStatusError(Exception):
+    """
+    Raised when a session is in a non-healthy state. Such as after a session failed or got expired.
+    """
+
+    def __init__(self, errors: list[SessionErrorData]):
+        message = f"Session is in an unhealthy state. Details: {[str(e) for e in errors]}"
+
+        super().__init__(self, message)
+
+
 class AuraApi:
     API_VERSION = "v1beta5"
 
@@ -126,7 +143,9 @@ def get_or_create_session(
         self._check_resp(response)
 
         raw_json: dict[str, Any] = response.json()
-        return SessionDetails.from_json(raw_json["data"], raw_json.get("errors", []))
+        self._check_errors(raw_json)
+
+        return SessionDetails.from_json(raw_json["data"])
 
     def get_session(self, session_id: str) -> Optional[SessionDetails]:
         response = self._request_session.get(
@@ -139,9 +158,11 @@ def get_session(self, session_id: str) -> Optional[SessionDetails]:
         self._check_resp(response)
 
         raw_json: dict[str, Any] = response.json()
-        return SessionDetails.from_json(raw_json["data"], raw_json.get("errors", []))
+        self._check_errors(raw_json)
+
+        return SessionDetails.from_json(raw_json["data"])
 
-    def list_sessions(self, dbid: Optional[str] = None) -> list[SessionDetails]:
+    def list_sessions(self, dbid: Optional[str] = None) -> list[SessionDetailsWithErrors]:
         # these are query parameters (not passed in the body)
         params = {
             "tenantId": self._tenant_id,
@@ -161,7 +182,7 @@ def list_sessions(self, dbid: Optional[str] = None) -> list[SessionDetails]:
         for error in raw_json.get("errors", []):
             errors_per_session[error["id"]].append(error)
 
-        return [SessionDetails.from_json(s, errors_per_session[s["id"]]) for s in data]
+        return [SessionDetailsWithErrors.from_json_with_error(s, errors_per_session[s["id"]]) for s in data]  # noqa: F821
 
     def wait_for_session_running(
         self,
@@ -177,14 +198,6 @@ def wait_for_session_running(
                 return WaitResult.from_error(f"Session `{session_id}` not found -- please retry")
             elif session.status == "Ready":
                 return WaitResult.from_connection_url(session.bolt_connection_url())
-            elif session.status == "Failed":
-                return WaitResult.from_error(
-                    f"Session `{session_id}` with name `{session.name}` failed due to: {session.errors}"
-                )
-            elif session.is_expired():
-                return WaitResult.from_error(
-                    f"Session `{session_id}` with name `{session.name}` is expired. Expired due to: {session.errors}"
-                )
             else:
                 self._logger.debug(
                     f"Session `{session_id}` is not yet running. "
@@ -324,6 +337,13 @@ def tenant_details(self) -> TenantDetails:
             self._tenant_details = TenantDetails.from_json(response.json()["data"])
         return self._tenant_details
 
+    def _check_errors(self, raw_json: dict[str, Any]) -> None:
+        errors = raw_json.get("errors", [])
+        typed_errors = [SessionErrorData.from_json(error) for error in errors] if errors else None
+
+        if typed_errors:
+            raise SessionStatusError(typed_errors)
+
     def _check_resp(self, resp: requests.Response) -> None:
         self._check_status_code(resp)
         self._check_endpoint_deprecation(resp)

graphdatascience/session/aura_api_responses.py

@@ -27,18 +27,15 @@ class SessionDetails:
     user_id: str
     tenant_id: str
     cloud_location: Optional[CloudLocation] = None
-    errors: Optional[list[SessionError]] = None
 
     @classmethod
-    def from_json(cls, data: dict[str, Any], errors: list[dict[str, Any]]) -> SessionDetails:
+    def from_json(cls, data: dict[str, Any]) -> SessionDetails:
         id = data["id"]
         expiry_date = data.get("expiry_date")
         ttl: Any | None = data.get("ttl")
         instance_id = data.get("instance_id")
         cloud_location = CloudLocation(data["cloud_provider"], data["region"]) if data.get("cloud_provider") else None
 
-        session_errors = [SessionError.from_json(error) for error in errors] if errors else None
-
         return cls(
             id=id,
             name=data["name"],
@@ -52,18 +49,45 @@ def from_json(cls, data: dict[str, Any], errors: list[dict[str, Any]]) -> Sessio
             tenant_id=data["tenant_id"],
             user_id=data["user_id"],
             cloud_location=cloud_location,
-            errors=session_errors,
         )
 
     def bolt_connection_url(self) -> str:
         return f"neo4j+s://{self.host}"
 
-    def is_expired(self) -> bool:
-        return self.status == "Expired"
+
+@dataclass(repr=True, frozen=True)
+class SessionDetailsWithErrors(SessionDetails):
+    errors: Optional[list[SessionErrorData]] = None
+
+    @classmethod
+    def from_json_with_error(cls, data: dict[str, Any], errors: list[dict[str, Any]]) -> SessionDetailsWithErrors:
+        session_errors = [SessionErrorData.from_json(error) for error in errors] if errors else None
+
+        id = data["id"]
+        expiry_date = data.get("expiry_date")
+        ttl: Any | None = data.get("ttl")
+        instance_id = data.get("instance_id")
+        cloud_location = CloudLocation(data["cloud_provider"], data["region"]) if data.get("cloud_provider") else None
+
+        return cls(
+            id=id,
+            name=data["name"],
+            instance_id=instance_id if instance_id else None,
+            memory=SessionMemoryValue.fromApiResponse(data["memory"]),
+            status=data["status"],
+            host=data["host"],
+            expiry_date=TimeParser.fromisoformat(expiry_date) if expiry_date else None,
+            created_at=TimeParser.fromisoformat(data["created_at"]),
+            ttl=Timedelta(ttl).to_pytimedelta() if ttl else None,  # datetime has no support for parsing timedelta
+            tenant_id=data["tenant_id"],
+            user_id=data["user_id"],
+            cloud_location=cloud_location,
+            errors=session_errors,
+        )
 
 
 @dataclass(repr=True, frozen=True)
-class SessionError:
+class SessionErrorData:
     """
     Represents information about a session errors.
     Indicates that session is in `Failed` state.
@@ -77,12 +101,15 @@ class SessionError:
     reason: str
 
     @classmethod
-    def from_json(cls, json: dict[str, Any]) -> SessionError:
+    def from_json(cls, json: dict[str, Any]) -> SessionErrorData:
         return cls(
             reason=json["reason"],
             message=json["message"],
         )
 
+    def __str__(self) -> str:
+        return f"Reason: {self.reason}, Message: {self.message}"
+
 
 @dataclass(repr=True, frozen=True)
 class InstanceDetails:

graphdatascience/session/dedicated_sessions.py

@@ -116,8 +116,7 @@ def delete(self, *, session_name: Optional[str] = None, session_id: Optional[str
         return False
 
     def list(self, dbid: Optional[str] = None) -> list[SessionInfo]:
-        sessions: list[SessionDetails] = self._aura_api.list_sessions(dbid)
-
+        sessions = self._aura_api.list_sessions(dbid)
         return [SessionInfo.from_session_details(i) for i in sessions]
 
     def _find_existing_session(self, session_name: str) -> Optional[SessionDetails]:

graphdatascience/session/gds_sessions.py

@@ -127,7 +127,7 @@ def delete(self, *, session_name: Optional[str] = None, session_id: Optional[str
 
     def list(self) -> list[SessionInfo]:
         """
-        Retrieves the list of GDS sessions visible by the user asscociated by the given api-credentials.
+        Retrieves the list of GDS sessions visible by the user associated by the given api-credentials.
 
         Returns:
             A list of SessionInfo objects representing the GDS sessions.

graphdatascience/session/session_info.py

@@ -4,7 +4,7 @@
 from datetime import datetime, timedelta
 from typing import Optional
 
-from graphdatascience.session.aura_api_responses import SessionDetails, SessionError
+from graphdatascience.session.aura_api_responses import SessionDetails, SessionDetailsWithErrors, SessionErrorData
 from graphdatascience.session.cloud_location import CloudLocation
 from graphdatascience.session.session_sizes import SessionMemoryValue
 
@@ -38,10 +38,14 @@ class SessionInfo:
     user_id: str
     cloud_location: Optional[CloudLocation]
     ttl: Optional[timedelta] = None
-    errors: Optional[list[SessionError]] = None
+    errors: Optional[list[SessionErrorData]] = None
 
     @classmethod
-    def from_session_details(cls, details: SessionDetails) -> SessionInfo:
+    def from_session_details(cls, details: SessionDetailsWithErrors | SessionDetails) -> SessionInfo:
+        errors: Optional[list[SessionErrorData]] = None
+        if isinstance(details, SessionDetailsWithErrors):
+            errors = details.errors
+
         return SessionInfo(
             id=details.id,
             name=details.name,
@@ -53,5 +57,5 @@ def from_session_details(cls, details: SessionDetails) -> SessionInfo:
             user_id=details.user_id,
             cloud_location=details.cloud_location,
             ttl=details.ttl,
-            errors=details.errors,
+            errors=errors,
         )