Merge pull request #682 from FlorentinD/fix-gds-version-init

Improve Aura-API error message / Refactor AuraAPI

Author: FlorentinD

graphdatascience/session/aura_api.py

@@ -3,11 +3,11 @@
 import logging
 import os
 import time
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 from urllib.parse import urlparse
 
-import requests as req
-from requests import HTTPError
+import requests
+import requests.auth
 
 from graphdatascience.session.algorithm_category import AlgorithmCategory
 from graphdatascience.session.aura_api_responses import (
@@ -23,21 +23,13 @@
 from graphdatascience.version import __version__
 
 
-class AuraApi:
-    class AuraAuthToken:
-        access_token: str
-        expires_in: int
-        token_type: str
-
-        def __init__(self, json: Dict[str, Any]) -> None:
-            self.access_token = json["access_token"]
-            expires_in: int = json["expires_in"]
-            self.expires_at = int(time.time()) + expires_in
-            self.token_type = json["token_type"]
+class AuraApiError(Exception):
+    def __init__(self, message: str, status_code: int):
+        super().__init__(self, message)
+        self.status_code = status_code
 
-        def is_expired(self) -> bool:
-            return self.expires_at >= int(time.time())
 
+class AuraApi:
     def __init__(self, client_id: str, client_secret: str, tenant_id: Optional[str] = None) -> None:
         self._dev_env = os.environ.get("AURA_ENV")
 
@@ -48,9 +40,13 @@ def __init__(self, client_id: str, client_secret: str, tenant_id: Optional[str]
         else:
             self._base_uri = f"https://api-{self._dev_env}.neo4j-dev.io"
 
-        self._credentials = (client_id, client_secret)
-        self._token: Optional[AuraApi.AuraAuthToken] = None
+        self._auth = AuraApi.Auth(oauth_url=f"{self._base_uri}/oauth/token", credentials=(client_id, client_secret))
         self._logger = logging.getLogger()
+
+        self._request_session = requests.Session()
+        self._request_session.headers = {"User-agent": f"neo4j-graphdatascience-v{__version__}"}
+        self._request_session.auth = self._auth
+
         self._tenant_id = tenant_id if tenant_id else self._get_tenant_id()
         self._tenant_details: Optional[TenantDetails] = None
 
@@ -64,36 +60,33 @@ 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 = self._request_session.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 = self._request_session.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 = self._request_session.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,9 +120,8 @@ def wait_for_session_running(
         )
 
     def delete_session(self, session_id: str, dbid: str) -> bool:
-        response = req.delete(
+        response = self._request_session.delete(
             f"{self._base_uri}/v1beta5/data-science/sessions/{session_id}",
-            headers=self._build_header(),
             json={"instance_id": dbid},
         )
 
@@ -138,7 +130,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 +149,38 @@ def create_instance(
             "cloud_provider": cloud_provider,
         }
 
-        response = req.post(
-            f"{self._base_uri}/v1/instances",
-            json=data,
-            headers=self._build_header(),
-        )
+        response = self._request_session.post(f"{self._base_uri}/v1/instances", json=data)
 
-        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(
-            f"{self._base_uri}/v1/instances/{instance_id}",
-            headers=self._build_header(),
-        )
+        response = self._request_session.delete(f"{self._base_uri}/v1/instances/{instance_id}")
 
         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(
-            f"{self._base_uri}/v1/instances",
-            headers=self._build_header(),
-            params={"tenantId": self._tenant_id},
-        )
+        response = self._request_session.get(f"{self._base_uri}/v1/instances", 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(
-            f"{self._base_uri}/v1/instances/{instance_id}",
-            headers=self._build_header(),
-        )
+        response = self._request_session.get(f"{self._base_uri}/v1/instances/{instance_id}")
 
         if response.status_code == 404:
             return None
 
-        response.raise_for_status()
+        self._check_code(response)
 
         raw_data = response.json()["data"]
 
@@ -246,17 +220,14 @@ 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 = self._request_session.post(f"{self._base_uri}/v1/instances/sizing", json=data)
+        self._check_code(response)
 
         return EstimationDetails.from_json(response.json()["data"])
 
     def _get_tenant_id(self) -> str:
-        response = req.get(
-            f"{self._base_uri}/v1/tenants",
-            headers=self._build_header(),
-        )
-        response.raise_for_status()
+        response = self._request_session.get(f"{self._base_uri}/v1/tenants")
+        self._check_code(response)
 
         raw_data = response.json()["data"]
 
@@ -270,36 +241,68 @@ def _get_tenant_id(self) -> str:
 
     def tenant_details(self) -> TenantDetails:
         if not self._tenant_details:
-            response = req.get(
-                f"{self._base_uri}/v1/tenants/{self._tenant_id}",
-                headers=self._build_header(),
-            )
-            response.raise_for_status()
+            response = self._request_session.get(f"{self._base_uri}/v1/tenants/{self._tenant_id}")
+            self._check_code(response)
             self._tenant_details = TenantDetails.from_json(response.json()["data"])
         return self._tenant_details
 
-    def _build_header(self) -> Dict[str, str]:
-        return {"Authorization": f"Bearer {self._auth_token()}", "User-agent": f"neo4j-graphdatascience-v{__version__}"}
-
-    def _auth_token(self) -> str:
-        if self._token is None or self._token.is_expired():
-            self._token = self._update_token()
-        return self._token.access_token
-
-    def _update_token(self) -> AuraAuthToken:
-        data = {
-            "grant_type": "client_credentials",
-        }
-
-        self._logger.debug("Updating oauth token")
-
-        response = req.post(
-            f"{self._base_uri}/oauth/token", data=data, auth=(self._credentials[0], self._credentials[1])
-        )
-
-        response.raise_for_status()
-
-        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"
+
+    class Auth(requests.auth.AuthBase):
+        class Token:
+            access_token: str
+            expires_in: int
+            token_type: str
+
+            def __init__(self, json: Dict[str, Any]) -> None:
+                self.access_token = json["access_token"]
+                self.token_type = json["token_type"]
+
+                expires_in: int = json["expires_in"]
+                refresh_in: int = expires_in if expires_in <= 10 else expires_in - 10
+                # avoid token expiry during request send by refreshing 10 seconds earlier
+                self.refresh_at = int(time.time()) + refresh_in
+
+            def should_refresh(self) -> bool:
+                return self.refresh_at >= int(time.time())
+
+        def __init__(self, oauth_url: str, credentials: Tuple[str, str]) -> None:
+            self._token: Optional[AuraApi.Auth.Token] = None
+            self._logger = logging.getLogger()
+            self._oauth_url = oauth_url
+            self._credentials = credentials
+
+        def __call__(self, r: requests.PreparedRequest) -> requests.PreparedRequest:
+            r.headers["Authorization"] = f"Bearer {self._auth_token()}"
+            return r
+
+        def _auth_token(self) -> str:
+            if self._token is None or self._token.should_refresh():
+                self._token = self._update_token()
+            return self._token.access_token
+
+        def _update_token(self) -> AuraApi.Auth.Token:
+            data = {
+                "grant_type": "client_credentials",
+            }
+
+            self._logger.debug("Updating oauth token")
+
+            resp = requests.post(self._oauth_url, data=data, auth=(self._credentials[0], self._credentials[1]))
+
+            if resp.status_code >= 400:
+                raise AuraApiError(
+                    "Failed to authorize with provided client credentials: "
+                    + f"{resp.status_code} - {resp.reason}, {resp.text}",
+                    status_code=resp.status_code,
+                )
+
+            return AuraApi.Auth.Token(resp.json())

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: