Fetch tenant details

for available regions and if tenant has ds capabilities

Author: FlorentinD

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, 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,35 @@ def from_json(cls, json: dict[str, Any]) -> InstanceCreateDetails:
         return cls(**{f.name: json[f.name] for f in fields})
 
 
+@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"])
+
+        if not ds_type:
+            raise RuntimeError(f"Tenant cannot create DS instances. Available instances are `{instance_types}`.")
+
+        return cls(
+            id=json["id"],
+            ds_type=ds_type,
+            regions_per_provider=regions_per_provider,
+        )
+
+
 class AuraApi:
     class AuraAuthToken:
         access_token: str
@@ -99,6 +129,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 +141,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,
         }
@@ -216,6 +247,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/tests/unit/test_aura_api.py

@@ -9,6 +9,7 @@
     AuraApi,
     InstanceCreateDetails,
     InstanceSpecificDetails,
+    TenantDetails,
 )
 
 
@@ -346,3 +347,36 @@ def test_parse_create_details() -> None:
     InstanceCreateDetails.from_json(
         {"id": "1", "username": "mats", "password": "1234", "connection_url": "url", "region": "fooo"}
     )
+
+
+def test_parse_tenant_details() -> None:
+    details = TenantDetails.from_json(
+        {
+            "id": "42",
+            "instance_configurations": [
+                {"type": "enterprise-db", "region": "eu-west1", "cloud_provider": "aws"},
+                {"type": "enterprise-ds", "region": "eu-west3", "cloud_provider": "gcp"},
+                {"type": "enterprise-ds", "region": "us-central1", "cloud_provider": "aws"},
+                {"type": "enterprise-ds", "region": "us-central3", "cloud_provider": "aws"},
+            ],
+        }
+    )
+
+    expected_details = TenantDetails(
+        "42", ds_type="enterprise-ds", regions_per_provider={"gcp": {"eu-west3"}, "aws": {"us-central1", "us-central3"}}
+    )
+    assert details == expected_details
+
+
+def test_parse_non_ds_details() -> None:
+    with pytest.raises(
+        RuntimeError, match="Tenant cannot create DS instances. Available instances are `{'enterprise-db'}`."
+    ):
+        TenantDetails.from_json(
+            {
+                "id": "42",
+                "instance_configurations": [
+                    {"type": "enterprise-db", "region": "europe-west1", "cloud_provider": "aws"}
+                ],
+            }
+        )