fix: prevent Client.logout from raising exceptions (#382)

Signed-off-by: Olivier Léobal <[email protected]>
Co-authored-by: Sarah Diot-Girard <[email protected]>

substra/sdk/backends/remote/rest_client.py

@@ -18,8 +18,8 @@
 logger = logging.getLogger(__name__)
 
 
-def _warn_of_session_expiration(expiration: str) -> None:
-    log_level = logging.INFO
+def _warn_of_session_expiration(expiration: str, minimum_log_level=logging.INFO) -> None:
+    log_level = minimum_log_level
     try:
         duration = datetime.fromisoformat(expiration) - datetime.now(timezone.utc)
         if duration.days:
@@ -28,7 +28,7 @@ def _warn_of_session_expiration(expiration: str) -> None:
             duration_msg = f"{duration.seconds//3600} hours"
         expires_at = f"in {duration_msg} ({expiration})"
         if duration < timedelta(hours=4):
-            log_level = logging.WARNING
+            log_level = max(log_level, logging.WARNING)
     except ValueError:  # Python 3.11 parses more formats than previous versions
         expires_at = expiration
     logger.log(log_level, f"Your session will expire {expires_at}")
@@ -52,7 +52,7 @@ def __init__(self, url, insecure, token):
         if not url:
             raise exceptions.SDKException("url required to connect to the Substra server")
         self._base_url = url[:-1] if url.endswith("/") else url
-        self._token_id = None  # filled in by self.login
+        self._token = None  # filled in by self.login
 
     def login(self, username, password):
         # we do not use self._headers in order to avoid existing tokens to be sent alongside the
@@ -70,28 +70,23 @@ def login(self, username, password):
         try:
             r = requests.post(f"{self._base_url}/api-token-auth/", data=data, headers=headers)
             r.raise_for_status()
-        except requests.exceptions.ConnectionError as e:
-            raise exceptions.ConnectionError.from_request_exception(e)
-        except requests.exceptions.Timeout as e:
-            raise exceptions.Timeout.from_request_exception(e)
-        except requests.exceptions.HTTPError as e:
-            logger.error(f"Requests error status {e.response.status_code}: {e.response.text}")
-
-            if e.response.status_code in (400, 401):
-                raise exceptions.BadLoginException.from_request_exception(e)
-
-            if (
-                e.response.status_code == 403
-                and getattr(e.response, "substra_identifier", None) == "implicit_login_disabled"
-            ):
-                raise exceptions.UsernamePasswordLoginDisabledException.from_request_exception(e)
-
-            raise exceptions.HTTPError.from_request_exception(e)
+        except requests.exceptions.RequestException as e:
+            if isinstance(e, requests.exceptions.HTTPError):
+                if e.response.status_code in (400, 401):
+                    raise exceptions.BadLoginException.from_request_exception(e)
+                if (
+                    e.response.status_code == 403
+                    and getattr(e.response, "substra_identifier", None) == "implicit_login_disabled"
+                ):
+                    raise exceptions.UsernamePasswordLoginDisabledException.from_request_exception(e)
+            raise exceptions.from_request_exception(e)
 
         try:
             rj = r.json()
+            if not all([it in rj for it in ["token", "expires_at", "id"]]):
+                raise exceptions.InvalidResponse(r, "Token object does not contain required fields")
             token = rj["token"]
-            self._token_id = rj["id"]
+            self._token = rj
         except json.decoder.JSONDecodeError:
             # sometimes requests seem to be fine, but the json is not being found
             # this might be if the url seems to be correct (in the syntax)
@@ -105,30 +100,34 @@ def login(self, username, password):
             "you should generate a token on the web UI instead: "
             "https://docs.substra.org/en/stable/documentation/api_tokens_generation.html"
         )
-        _warn_of_session_expiration(rj["expires_at"])
+        _warn_of_session_expiration(self._token["expires_at"])
         self._headers["Authorization"] = f"Token {token}"
 
         return token
 
     def logout(self) -> None:
-        if self._token_id is None:
+        if self._token is None:
             return
         try:
             r = requests.delete(
-                f"{self._base_url}/active-api-tokens/", params={"id": self._token_id}, headers=self._headers
+                f"{self._base_url}/active-api-tokens/", params={"id": self._token["id"]}, headers=self._headers
             )
             r.raise_for_status()
-            self._token_id = None
+            self._token = None
             logger.info("Successfully logged out")
-        except requests.exceptions.ConnectionError as e:
-            raise exceptions.ConnectionError.from_request_exception(e)
-        except requests.exceptions.Timeout as e:
-            raise exceptions.Timeout.from_request_exception(e)
-        except requests.exceptions.HTTPError as e:
-            raise exceptions.HTTPError.from_request_exception(e)
-
-    # TODO: '__request' is too complex, consider refactoring
-    def __request(self, request_name, url, **request_kwargs):  # noqa: C901
+        except requests.exceptions.RequestException as e:
+            logger.error(
+                f"Could not end session {self._token['id']}"
+                + (
+                    f", got {e.response.status_code}: {e.response.text}"
+                    if isinstance(e, requests.exceptions.HTTPError)
+                    else f": {e}"
+                )
+            )
+            _warn_of_session_expiration(self._token["expires_at"], minimum_log_level=logging.WARNING)
+            # this isn't too much of an issue, the token will expire on its own
+
+    def __request(self, request_name, url, **request_kwargs):
         """Base request helper."""
 
         if request_name == "get":
@@ -154,41 +153,8 @@ def __request(self, request_name, url, **request_kwargs):  # noqa: C901
             r = fn(url, headers=self._headers, **kwargs)
             r.raise_for_status()
 
-        except requests.exceptions.ConnectionError as e:
-            raise exceptions.ConnectionError.from_request_exception(e)
-
-        except requests.exceptions.Timeout as e:
-            raise exceptions.Timeout.from_request_exception(e)
-
-        except requests.exceptions.HTTPError as e:
-            logger.error(f"Requests error status {e.response.status_code}: {e.response.text}")
-
-            if e.response.status_code == 400:
-                raise exceptions.InvalidRequest.from_request_exception(e)
-
-            if e.response.status_code == 401:
-                raise exceptions.AuthenticationError.from_request_exception(e)
-
-            if e.response.status_code == 403:
-                raise exceptions.AuthorizationError.from_request_exception(e)
-
-            if e.response.status_code == 404:
-                raise exceptions.NotFound.from_request_exception(e)
-
-            if e.response.status_code == 408:
-                raise exceptions.RequestTimeout.from_request_exception(e)
-
-            if e.response.status_code == 409:
-                raise exceptions.AlreadyExists.from_request_exception(e)
-
-            if e.response.status_code == 500:
-                raise exceptions.InternalServerError.from_request_exception(e)
-
-            if e.response.status_code in [502, 503, 504]:
-                raise exceptions.GatewayUnavailable.from_request_exception(e)
-
-            raise exceptions.HTTPError.from_request_exception(e)
-
+        except requests.exceptions.RequestException as e:
+            raise exceptions.from_request_exception(e)
         return r
 
     @utils.retry_on_exception(exceptions=(exceptions.GatewayUnavailable))

substra/sdk/exceptions.py

@@ -1,3 +1,11 @@
+import logging
+from typing import Union
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
 class SDKException(Exception):
     pass
 
@@ -160,6 +168,39 @@ def from_request_exception(cls, request_exception):
         )
 
 
+def from_request_exception(
+    e: requests.exceptions.RequestException,
+) -> Union[RequestException, requests.exceptions.RequestException]:
+    """
+    try turning an exception from the `requests` library into a Substra exception
+    """
+    connection_error_mapping: dict[requests.exceptions.RequestException, RequestException] = {
+        requests.exceptions.ConnectionError: ConnectionError,
+        requests.exceptions.Timeout: Timeout,
+    }
+    for k, v in connection_error_mapping.items():
+        if isinstance(e, k):
+            return v.from_request_exception(e)
+
+    http_status_mapping: dict[int, RequestException] = {
+        400: InvalidRequest,
+        401: AuthenticationError,
+        403: AuthorizationError,
+        404: NotFound,
+        408: RequestTimeout,
+        409: AlreadyExists,
+        500: InternalServerError,
+        502: GatewayUnavailable,
+        503: GatewayUnavailable,
+        504: GatewayUnavailable,
+    }
+    if isinstance(e, requests.exceptions.HTTPError):
+        logger.error(f"Requests error status {e.response.status_code}: {e.response.text}")
+        return http_status_mapping.get(e.response.status_code, HTTPError).from_request_exception(e)
+
+    return e
+
+
 class ConfigurationInfoError(SDKException):
     """ConfigurationInfoError"""