Merge pull request #584 from FlorentinD/try-sessions-against-production

Validate create requests against tenants available configurations

Author: Mats-SX

examples/dev/gds-sessions.ipynb

@@ -70,18 +70,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from graphdatascience.gds_session.dbms_connection_info import DbmsConnectionInfo\n",
+    "from graphdatascience.gds_session import DbmsConnectionInfo\n",
     "import os\n",
     "\n",
     "from neo4j import GraphDatabase\n",
     "\n",
-    "# We need to tell the GDS client that we are working with a development environment.\n",
-    "# This does not need to be set in production.\n",
-    "os.environ[\"AURA_ENV\"] = \"devstrawberryfield\"\n",
+    "# We can tell the GDS client that we are working with a development environment.\n",
+    "# os.environ[\"AURA_ENV\"] = \"devstrawberryfield\"\n",
     "\n",
-    "db_connection_info = DbmsConnectionInfo(\n",
-    "    f\"neo4j+s://{db_id}-{os.environ['AURA_ENV']}.databases.neo4j-dev.io\", \"neo4j\", db_password\n",
+    "uri = (\n",
+    "    f\"neo4j+s://{db_id}-{os.environ['AURA_ENV']}.databases.neo4j-dev.io\"\n",
+    "    if os.environ.get(\"AURA_ENV\")\n",
+    "    else f\"neo4j+s://{db_id}.databases.neo4j.io\"\n",
     ")\n",
+    "\n",
+    "db_connection_info = DbmsConnectionInfo(uri, \"neo4j\", db_password)\n",
     "# start a standard Neo4j Python Driver to connect to the AuraDB instance\n",
     "driver = GraphDatabase.driver(db_connection_info.uri, auth=db_connection_info.auth())\n",
     "\n",
@@ -171,9 +174,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Let's call this function and see what it returns\n",
+    "# Let's call this function and verify the arrow server is enabled\n",
     "with driver.session() as session:\n",
-    "    display(session.run(\"CALL internal.arrow.status\").to_df())"
+    "    arrow_status = session.run(\"CALL internal.arrow.status\").to_df()\n",
+    "    print(arrow_status)\n",
+    "    assert arrow_status[\"enabled\"][0] == True, \"Arrow server on the db needs to be enabled\""
    ]
   },
   {
@@ -244,7 +249,7 @@
    "outputs": [],
    "source": [
     "# The new stuff!\n",
-    "from graphdatascience.gds_session.gds_sessions import GdsSessions, AuraAPICredentials\n",
+    "from graphdatascience.gds_session import GdsSessions, AuraAPICredentials\n",
     "\n",
     "# Create a new AuraSessions object\n",
     "sessions = GdsSessions(ds_connection=AuraAPICredentials(CLIENT_ID, CLIENT_SECRET))"

graphdatascience/gds_session/__init__.py

@@ -0,0 +1,11 @@
+from .dbms_connection_info import DbmsConnectionInfo
+from .gds_sessions import AuraAPICredentials, GdsSessions, SessionInfo
+from .session_sizes import SessionSizes
+
+__all__ = [
+    "GdsSessions",
+    "SessionInfo",
+    "DbmsConnectionInfo",
+    "AuraAPICredentials",
+    "SessionSizes",
+]

graphdatascience/gds_session/aura_api.py

@@ -4,8 +4,9 @@
 import logging
 import os
 import time
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, List, Optional
+from typing import Any, List, NamedTuple, Optional, Set
 from urllib.parse import urlparse
 
 import requests as req
@@ -14,7 +15,7 @@
 from graphdatascience.version import __version__
 
 
-@dataclass(repr=True)
+@dataclass(repr=True, frozen=True)
 class InstanceDetails:
     id: str
     name: str
@@ -31,7 +32,7 @@ def fromJson(cls, json: dict[str, Any]) -> InstanceDetails:
         )
 
 
-@dataclass(repr=True)
+@dataclass(repr=True, frozen=True)
 class InstanceSpecificDetails(InstanceDetails):
     status: str
     connection_url: str
@@ -54,7 +55,7 @@ def fromJson(cls, json: dict[str, Any]) -> InstanceSpecificDetails:
         )
 
 
-@dataclass(repr=True)
+@dataclass(repr=True, frozen=True)
 class InstanceCreateDetails:
     id: str
     username: str
@@ -70,6 +71,51 @@ def from_json(cls, json: dict[str, Any]) -> InstanceCreateDetails:
         return cls(**{f.name: json[f.name] for f in fields})
 
 
+class WaitResult(NamedTuple):
+    connection_url: str
+    error: str
+
+    @classmethod
+    def from_error(cls, error: str) -> WaitResult:
+        return cls(connection_url="", error=error)
+
+    @classmethod
+    def from_connection_url(cls, connection_url: str) -> WaitResult:
+        return cls(connection_url=connection_url, error="")
+
+
+@dataclass(repr=True, frozen=True)
+class TenantDetails:
+    id: str
+    ds_type: str
+    regions_per_provider: dict[str, Set[str]]
+
+    @classmethod
+    def from_json(cls, json: dict[str, Any]) -> TenantDetails:
+        regions_per_provider = defaultdict(set)
+        instance_types = set()
+        ds_type = None
+
+        for configs in json["instance_configurations"]:
+            type = configs["type"]
+            if type.split("-")[1] == "ds":
+                regions_per_provider[configs["cloud_provider"]].add(configs["region"])
+                ds_type = type
+            instance_types.add(configs["type"])
+
+        id = json["id"]
+        if not ds_type:
+            raise RuntimeError(
+                f"Tenant with id `{id}` cannot create DS instances. Available instances are `{instance_types}`."
+            )
+
+        return cls(
+            id=id,
+            ds_type=ds_type,
+            regions_per_provider=regions_per_provider,
+        )
+
+
 class AuraApi:
     class AuraAuthToken:
         access_token: str
@@ -99,6 +145,7 @@ def __init__(self, client_id: str, client_secret: str, tenant_id: Optional[str]
         self._token: Optional[AuraApi.AuraAuthToken] = None
         self._logger = logging.getLogger()
         self._tenant_id = tenant_id if tenant_id else self._get_tenant_id()
+        self._tenant_details: Optional[TenantDetails] = None
 
     @staticmethod
     def extract_id(uri: str) -> str:
@@ -110,14 +157,14 @@ def extract_id(uri: str) -> str:
         return host.split(".")[0].split("-")[0]
 
     def create_instance(self, name: str, memory: str, cloud_provider: str, region: str) -> InstanceCreateDetails:
-        # TODO should give more control here
+        tenant_details = self.tenant_details()
+
         data = {
             "name": name,
             "memory": memory,
             "version": "5",
             "region": region,
-            # TODO should be figured out from the tenant details in the future
-            "type": self._instance_type(),
+            "type": tenant_details.ds_type,
             "tenant_id": self._tenant_id,
             "cloud_provider": cloud_provider,
         }
@@ -179,16 +226,16 @@ def list_instance(self, instance_id: str) -> Optional[InstanceSpecificDetails]:
 
     def wait_for_instance_running(
         self, instance_id: str, sleep_time: float = 0.2, max_sleep_time: float = 300
-    ) -> Optional[str]:
+    ) -> WaitResult:
         waited_time = 0.0
         while waited_time <= max_sleep_time:
             instance = self.list_instance(instance_id)
             if instance is None:
-                return "Instance is not found -- please retry"
+                return WaitResult.from_error("Instance is not found -- please retry")
             elif instance.status in ["deleting", "destroying"]:
-                return "Instance is being deleted"
+                return WaitResult.from_error("Instance is being deleted")
             elif instance.status == "running":
-                return None
+                return WaitResult.from_connection_url(instance.connection_url)
             else:
                 self._logger.debug(
                     f"Instance `{instance_id}` is not yet running. "
@@ -198,7 +245,7 @@ def wait_for_instance_running(
             waited_time += sleep_time
             time.sleep(sleep_time)
 
-        return f"Instance is not running after waiting for {waited_time} seconds"
+        return WaitResult.from_error(f"Instance is not running after waiting for {waited_time} seconds")
 
     def _get_tenant_id(self) -> str:
         response = req.get(
@@ -216,6 +263,16 @@ def _get_tenant_id(self) -> str:
 
         return raw_data[0]["id"]  # type: ignore
 
+    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()
+            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__}"}
 

graphdatascience/gds_session/gds_sessions.py

@@ -12,6 +12,7 @@
 )
 from graphdatascience.gds_session.aura_graph_data_science import AuraGraphDataScience
 from graphdatascience.gds_session.dbms_connection_info import DbmsConnectionInfo
+from graphdatascience.gds_session.region_suggester import closest_match
 from graphdatascience.gds_session.session_sizes import SessionSizeByMemory
 
 
@@ -45,7 +46,10 @@ def __init__(self, ds_connection: AuraAPICredentials) -> None:
         )
 
     def get_or_create(
-        self, session_name: str, size: SessionSizeByMemory, db_connection: DbmsConnectionInfo
+        self,
+        session_name: str,
+        size: SessionSizeByMemory,
+        db_connection: DbmsConnectionInfo,
     ) -> AuraGraphDataScience:
         connected_instance = self._try_connect(session_name, db_connection)
         if connected_instance is not None:
@@ -56,15 +60,17 @@ def get_or_create(
         if not db_instance:
             raise ValueError(f"Could not find Aura instance with the uri `{db_connection.uri}`")
 
+        region = self._ds_region(db_instance.region, db_instance.cloud_provider)
+
         create_details = self._aura_api.create_instance(
-            GdsSessions._instance_name(session_name), size.value, db_instance.cloud_provider, db_instance.region
+            GdsSessions._instance_name(session_name), size.value, db_instance.cloud_provider, region
         )
         wait_result = self._aura_api.wait_for_instance_running(create_details.id)
-        if wait_result is not None:
-            raise RuntimeError(f"Failed to create session `{session_name}`: {wait_result}")
+        if err := wait_result.error:
+            raise RuntimeError(f"Failed to create session `{session_name}`: {err}")
 
         gds_user = create_details.username
-        gds_url = create_details.connection_url
+        gds_url = wait_result.connection_url
 
         self._change_initial_pw(
             gds_url=gds_url, gds_user=gds_user, initial_pw=create_details.password, new_pw=db_connection.password
@@ -118,16 +124,24 @@ def _try_connect(self, session_name: str, db_connection: DbmsConnectionInfo) ->
         if len(matched_instances) > 1:
             self._fail_ambiguous_session(session_name, matched_instances)
 
-        instance_details = self._aura_api.list_instance(matched_instances[0].id)
+        wait_result = self._aura_api.wait_for_instance_running(matched_instances[0].id)
+        if err := wait_result.error:
+            raise RuntimeError(f"Failed to connect to session `{session_name}`: {err}")
+        gds_url = wait_result.connection_url
+
+        return self._construct_client(session_name=session_name, gds_url=gds_url, db_connection=db_connection)
 
-        if instance_details:
-            gds_url = instance_details.connection_url
-        else:
-            raise RuntimeError(
-                f"Unable to get connection information for session `{session_name}`. Does it still exist?"
+    def _ds_region(self, region: str, cloud_provider: str) -> str:
+        tenant_details = self._aura_api.tenant_details()
+        available_regions = tenant_details.regions_per_provider[cloud_provider]
+
+        match = closest_match(region, available_regions)
+        if not match:
+            raise ValueError(
+                f"Tenant `{tenant_details.id}` cannot create GDS sessions at cloud provider `{cloud_provider}`."
             )
 
-        return self._construct_client(session_name=session_name, gds_url=gds_url, db_connection=db_connection)
+        return match
 
     def _construct_client(
         self, session_name: str, gds_url: str, db_connection: DbmsConnectionInfo

graphdatascience/gds_session/region_suggester.py

@@ -0,0 +1,17 @@
+from typing import Iterable, Optional
+
+import textdistance
+
+
+# AuraDB and AuraDS regions are not the same, so we need to find the closest match.
+def closest_match(db_region: str, ds_regions: Iterable[str]) -> Optional[str]:
+    curr_max_similarity = 0.0
+    closest_option = None
+
+    for region in ds_regions:
+        similarity = textdistance.jaro_winkler(db_region, region)
+        if similarity > curr_max_similarity:
+            closest_option = region
+            curr_max_similarity = similarity
+
+    return closest_option