Throw error containing response text

Author: FlorentinD

graphdatascience/session/aura_api.py

@@ -6,8 +6,7 @@
 from typing import Any, Dict, List, Optional
 from urllib.parse import urlparse
 
-import requests as req
-from requests import HTTPError
+import requests
 
 from graphdatascience.session.algorithm_category import AlgorithmCategory
 from graphdatascience.session.aura_api_responses import (
@@ -22,6 +21,10 @@
 from graphdatascience.session.session_sizes import SessionMemoryValue
 from graphdatascience.version import __version__
 
+class AuraApiError(Exception):
+    def __init__(self, message: str, status_code: int):
+        super().__init__(self, message)
+        self.status_code = status_code
 
 class AuraApi:
     class AuraAuthToken:
@@ -64,36 +67,36 @@ def extract_id(uri: str) -> str:
         return host.split(".")[0].split("-")[0]
 
     def create_session(self, name: str, dbid: str, pwd: str, memory: SessionMemoryValue) -> SessionDetails:
-        response = req.post(
+        response = requests.post(
             f"{self._base_uri}/v1beta5/data-science/sessions",
             headers=self._build_header(),
             json={"name": name, "instance_id": dbid, "password": pwd, "memory": memory.value},
         )
 
-        response.raise_for_status()
+        self._check_code(response)
 
         return SessionDetails.fromJson(response.json())
 
     def list_session(self, session_id: str, dbid: str) -> Optional[SessionDetails]:
-        response = req.get(
+        response = requests.get(
             f"{self._base_uri}/v1beta5/data-science/sessions/{session_id}?instanceId={dbid}",
             headers=self._build_header(),
         )
 
         if response.status_code == 404:
             return None
 
-        response.raise_for_status()
+        self._check_code(response)
 
         return SessionDetails.fromJson(response.json())
 
     def list_sessions(self, dbid: str) -> List[SessionDetails]:
-        response = req.get(
+        response = requests.get(
             f"{self._base_uri}/v1beta5/data-science/sessions?instanceId={dbid}",
             headers=self._build_header(),
         )
 
-        response.raise_for_status()
+        self._check_code(response)
 
         return [SessionDetails.fromJson(s) for s in response.json()]
 
@@ -127,7 +130,7 @@ def wait_for_session_running(
         )
 
     def delete_session(self, session_id: str, dbid: str) -> bool:
-        response = req.delete(
+        response = requests.delete(
             f"{self._base_uri}/v1beta5/data-science/sessions/{session_id}",
             headers=self._build_header(),
             json={"instance_id": dbid},
@@ -138,7 +141,7 @@ def delete_session(self, session_id: str, dbid: str) -> bool:
         elif response.status_code == 202:
             return True
 
-        response.raise_for_status()
+        self._check_code(response)
 
         return False
 
@@ -157,56 +160,52 @@ def create_instance(
             "cloud_provider": cloud_provider,
         }
 
-        response = req.post(
+        response = requests.post(
             f"{self._base_uri}/v1/instances",
             json=data,
             headers=self._build_header(),
         )
 
-        try:
-            response.raise_for_status()
-        except HTTPError as e:
-            print(response.json())
-            raise e
+        self._check_code(response)
 
         return InstanceCreateDetails.from_json(response.json()["data"])
 
     def delete_instance(self, instance_id: str) -> Optional[InstanceSpecificDetails]:
-        response = req.delete(
+        response = requests.delete(
             f"{self._base_uri}/v1/instances/{instance_id}",
             headers=self._build_header(),
         )
 
         if response.status_code == 404:
             return None
 
-        response.raise_for_status()
+        self._check_code(response)
 
         return InstanceSpecificDetails.fromJson(response.json()["data"])
 
     def list_instances(self) -> List[InstanceDetails]:
-        response = req.get(
+        response = requests.get(
             f"{self._base_uri}/v1/instances",
             headers=self._build_header(),
             params={"tenantId": self._tenant_id},
         )
 
-        response.raise_for_status()
+        self._check_code(response)
 
         raw_data = response.json()["data"]
 
         return [InstanceDetails.fromJson(i) for i in raw_data]
 
     def list_instance(self, instance_id: str) -> Optional[InstanceSpecificDetails]:
-        response = req.get(
+        response = requests.get(
             f"{self._base_uri}/v1/instances/{instance_id}",
             headers=self._build_header(),
         )
 
         if response.status_code == 404:
             return None
 
-        response.raise_for_status()
+        self._check_code(response)
 
         raw_data = response.json()["data"]
 
@@ -246,17 +245,17 @@ def estimate_size(
             "instance_type": "dsenterprise",
         }
 
-        response = req.post(f"{self._base_uri}/v1/instances/sizing", headers=self._build_header(), json=data)
-        response.raise_for_status()
+        response = requests.post(f"{self._base_uri}/v1/instances/sizing", headers=self._build_header(), json=data)
+        self._check_code(response)
 
         return EstimationDetails.from_json(response.json()["data"])
 
     def _get_tenant_id(self) -> str:
-        response = req.get(
+        response = requests.get(
             f"{self._base_uri}/v1/tenants",
             headers=self._build_header(),
         )
-        response.raise_for_status()
+        self._check_code(response)
 
         raw_data = response.json()["data"]
 
@@ -270,11 +269,11 @@ def _get_tenant_id(self) -> str:
 
     def tenant_details(self) -> TenantDetails:
         if not self._tenant_details:
-            response = req.get(
+            response = requests.get(
                 f"{self._base_uri}/v1/tenants/{self._tenant_id}",
                 headers=self._build_header(),
             )
-            response.raise_for_status()
+            self._check_code(response)
             self._tenant_details = TenantDetails.from_json(response.json()["data"])
         return self._tenant_details
 
@@ -293,13 +292,20 @@ def _update_token(self) -> AuraAuthToken:
 
         self._logger.debug("Updating oauth token")
 
-        response = req.post(
+        response = requests.post(
             f"{self._base_uri}/oauth/token", data=data, auth=(self._credentials[0], self._credentials[1])
         )
 
-        response.raise_for_status()
+        self._check_code(response)
 
         return AuraApi.AuraAuthToken(response.json())
 
+    def _check_code(self, resp: requests.Response) -> None:
+        if resp.status_code >= 400:
+            raise AuraApiError(
+                f"Request for {resp.url} failed with status code {resp.status_code} - {resp.reason}: {resp.text}",
+                status_code=resp.status_code,
+            )
+
     def _instance_type(self) -> str:
         return "enterprise-ds" if not self._dev_env else "professional-ds"

graphdatascience/session/dedicated_sessions.py

@@ -5,10 +5,8 @@
 from datetime import datetime, timedelta, timezone
 from typing import List, Optional
 
-from requests.exceptions import HTTPError
-
 from graphdatascience.session.algorithm_category import AlgorithmCategory
-from graphdatascience.session.aura_api import AuraApi
+from graphdatascience.session.aura_api import AuraApi, AuraApiError
 from graphdatascience.session.aura_api_responses import SessionDetails
 from graphdatascience.session.aura_graph_data_science import AuraGraphDataScience
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
@@ -97,9 +95,9 @@ def list(self) -> List[SessionInfo]:
         for db in dbs:
             try:
                 sessions.extend(self._aura_api.list_sessions(db.id))
-            except HTTPError as e:
+            except AuraApiError as e:
                 # ignore 404 errors when listing sessions as it could mean paused sessions or deleted sessions
-                if e.response.status_code != 404:
+                if e.status_code != 404:
                     raise e
 
         return [SessionInfo.from_session_details(i) for i in sessions]
@@ -108,9 +106,9 @@ def _find_existing_session(self, session_name: str, dbid: str) -> Optional[Sessi
         matched_sessions: List[SessionDetails] = []
         try:
             matched_sessions = [s for s in self._aura_api.list_sessions(dbid) if s.name == session_name]
-        except HTTPError as e:
+        except AuraApiError as e:
             # ignore 404 errors when listing sessions as it could mean paused sessions or deleted sessions
-            if e.response.status_code != 404:
+            if e.status_code != 404:
                 raise e
 
         if len(matched_sessions) == 0:

graphdatascience/tests/unit/test_dedicated_sessions.py

@@ -9,7 +9,7 @@
 from requests.exceptions import HTTPError
 
 from graphdatascience.session.algorithm_category import AlgorithmCategory
-from graphdatascience.session.aura_api import AuraApi
+from graphdatascience.session.aura_api import AuraApi, AuraApiError
 from graphdatascience.session.aura_api_responses import (
     EstimationDetails,
     InstanceCreateDetails,
@@ -73,10 +73,7 @@ def add_session(self, session: SessionDetails) -> None:
     def _mimic_paused_db_behaviour(self, dbid: str) -> None:
         db = self.list_instance(dbid)
         if db and db.status == "paused":
-            response = Response()
-            response.status_code = 404
-            response._content = b"database not found"
-            raise HTTPError(request=None, response=response)
+            raise AuraApiError(message="Database not found", status_code=404)
 
     def create_instance(
         self, name: str, memory: SessionMemoryValue, cloud_provider: str, region: str