graphdatascience 1.9__tar.gz → 1.10a1__tar.gz

{graphdatascience-1.9/graphdatascience.egg-info → graphdatascience-1.10a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphdatascience
-Version: 1.9
+Version: 1.10a1
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -20,6 +20,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Database
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Software Development

{graphdatascience-1.9 → graphdatascience-1.10a1}/graphdatascience/__init__.py

@@ -1,4 +1,3 @@
-from .gds_session.gds_sessions import GdsSessions
 from .graph.graph_create_result import GraphCreateResult
 from .graph.graph_object import Graph
 from .graph_data_science import GraphDataScience
@@ -13,6 +12,7 @@ from .pipeline.nc_training_pipeline import NCTrainingPipeline
 from .pipeline.nr_training_pipeline import NRTrainingPipeline
 from .query_runner.query_runner import QueryRunner
 from .server_version.server_version import ServerVersion
+from .session.gds_sessions import GdsSessions
 from .version import __version__
 
 __all__ = [

{graphdatascience-1.9 → graphdatascience-1.10a1}/graphdatascience/error/endpoint_suggester.py

@@ -9,7 +9,7 @@ def generate_suggestive_error_message(requested_endpoint: str, all_endpoints: Li
     MIN_SIMILARITY_FOR_SUGGESTION = 0.9
 
     closest_endpoint = None
-    curr_max_similarity = 0
+    curr_max_similarity = 0.0
     for ep in all_endpoints:
         similarity = textdistance.jaro_winkler(requested_endpoint, ep)
         if similarity >= MIN_SIMILARITY_FOR_SUGGESTION:

graphdatascience-1.9/graphdatascience/graph/graph_proc_runner.py → graphdatascience-1.10a1/graphdatascience/graph/base_graph_proc_runner.py

@@ -1,6 +1,7 @@
 import os
 import pathlib
 import sys
+import warnings
 from typing import Any, Dict, List, Optional, Union
 
 import pandas as pd
@@ -24,7 +25,6 @@ from .graph_entity_ops_runner import (
 )
 from .graph_export_runner import GraphExportRunner
 from .graph_object import Graph
-from .graph_project_runner import GraphProjectRemoteRunner, GraphProjectRunner
 from .graph_sample_runner import GraphSampleRunner
 from .graph_type_check import (
     from_graph_type_check,
@@ -34,7 +34,6 @@ from .graph_type_check import (
 from .ogb_loader import OGBLLoader, OGBNLoader
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.graph.graph_create_result import GraphCreateResult
-from graphdatascience.graph.graph_cypher_runner import GraphCypherRunner
 
 Strings = Union[str, List[str]]
 
@@ -42,6 +41,15 @@ is_neo4j_4_driver = ServerVersion.from_string(neo4j_driver_version) < ServerVers
 
 
 class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
+    def __init__(self, query_runner: Any, namespace: str, server_version: ServerVersion):
+        super().__init__(query_runner, namespace, server_version)
+        # Pandas 2.2.0 deprecated an API used by ArrowTable.to_pandas() (< pyarrow 14.0)
+        warnings.filterwarnings(
+            "ignore",
+            category=DeprecationWarning,
+            message=r"Passing a BlockManager to DataFrame is deprecated",
+        )
+
     @staticmethod
     def _path(package: str, resource: str) -> pathlib.Path:
         if sys.version_info >= (3, 9):
@@ -558,22 +566,3 @@ class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
             endpoint=self._namespace,
             params=params,
         ).squeeze()
-
-
-class GraphProcRunner(BaseGraphProcRunner):
-    @property
-    def project(self) -> GraphProjectRunner:
-        self._namespace += ".project"
-        return GraphProjectRunner(self._query_runner, self._namespace, self._server_version)
-
-    @property
-    def cypher(self) -> GraphCypherRunner:
-        self._namespace += ".project"
-        return GraphCypherRunner(self._query_runner, self._namespace, self._server_version)
-
-
-class GraphRemoteProcRunner(BaseGraphProcRunner):
-    @property
-    def project(self) -> GraphProjectRemoteRunner:
-        self._namespace += ".project.remoteDb"
-        return GraphProjectRemoteRunner(self._query_runner, self._namespace, self._server_version)

{graphdatascience-1.9 → graphdatascience-1.10a1}/graphdatascience/graph/graph_entity_ops_runner.py

@@ -1,5 +1,6 @@
 from functools import reduce
 from typing import Any, Dict, List, Type, Union
+from warnings import filterwarnings
 
 import pandas as pd
 from pandas import DataFrame, Series
@@ -26,6 +27,13 @@ class TopologyDataFrame(DataFrame):
         return TopologyDataFrame
 
     def by_rel_type(self) -> Dict[str, List[List[int]]]:
+        # Pandas 2.2.0 deprecated an internal API used by DF.take(indices)
+        filterwarnings(
+            "ignore",
+            category=DeprecationWarning,
+            message=r"Passing a BlockManager to TopologyDataFrame is deprecated",
+        )
+
         gb = self.groupby("relationshipType", observed=True)
 
         output = {}
@@ -242,13 +250,13 @@ class ToUndirectedRunner(IllegalAttrChecker):
 
     @graph_type_check
     def __call__(self, G: Graph, relationship_type: str, mutate_relationship_type: str, **config: Any) -> "Series[Any]":
-        return self._run_procedure(G, relationship_type, mutate_relationship_type)
+        return self._run_procedure(G, relationship_type, mutate_relationship_type, **config)
 
     @graph_type_check
     @compatible_with("estimate", min_inclusive=ServerVersion(2, 3, 0))
     def estimate(self, G: Graph, relationship_type: str, mutate_relationship_type: str, **config: Any) -> "Series[Any]":
         self._namespace += ".estimate"
-        return self._run_procedure(G, relationship_type, mutate_relationship_type)
+        return self._run_procedure(G, relationship_type, mutate_relationship_type, **config)
 
 
 class GraphRelationshipsRunner(GraphEntityOpsBaseRunner):

graphdatascience-1.10a1/graphdatascience/graph/graph_proc_runner.py

@@ -0,0 +1,15 @@
+from .graph_project_runner import GraphProjectRunner
+from graphdatascience.graph.base_graph_proc_runner import BaseGraphProcRunner
+from graphdatascience.graph.graph_cypher_runner import GraphCypherRunner
+
+
+class GraphProcRunner(BaseGraphProcRunner):
+    @property
+    def project(self) -> GraphProjectRunner:
+        self._namespace += ".project"
+        return GraphProjectRunner(self._query_runner, self._namespace, self._server_version)
+
+    @property
+    def cypher(self) -> GraphCypherRunner:
+        self._namespace += ".project"
+        return GraphCypherRunner(self._query_runner, self._namespace, self._server_version)

{graphdatascience-1.9 → graphdatascience-1.10a1}/graphdatascience/graph/graph_project_runner.py

@@ -5,13 +5,10 @@ from typing import Any
 from pandas import Series
 
 from ..error.illegal_attr_checker import IllegalAttrChecker
-from ..gds_session.schema import NODE_PROPERTY_SCHEMA, RELATIONSHIP_PROPERTY_SCHEMA
 from .graph_object import Graph
 from .graph_type_check import from_graph_type_check
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.graph.graph_create_result import GraphCreateResult
-from graphdatascience.server_version.compatible_with import compatible_with
-from graphdatascience.server_version.server_version import ServerVersion
 
 
 class GraphProjectRunner(IllegalAttrChecker):
@@ -73,29 +70,3 @@ class GraphProjectBetaRunner(IllegalAttrChecker):
         ).squeeze()
 
         return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
-
-
-class GraphProjectRemoteRunner(IllegalAttrChecker):
-    @compatible_with("project", min_inclusive=ServerVersion(2, 6, 0))
-    def __call__(self, graph_name: str, query: str, **config: Any) -> GraphCreateResult:
-        placeholder = "<>"  # host and token will be added by query runner
-        self.map_property_types(config)
-        params = CallParameters(
-            graph_name=graph_name,
-            query=query,
-            token=placeholder,
-            host=placeholder,
-            remote_database=self._query_runner.database(),
-            config=config,
-        )
-        result = self._query_runner.call_procedure(
-            endpoint=self._namespace,
-            params=params,
-        ).squeeze()
-        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
-
-    @staticmethod
-    def map_property_types(config: dict[str, Any]) -> None:
-        for key in [NODE_PROPERTY_SCHEMA, RELATIONSHIP_PROPERTY_SCHEMA]:
-            if key in config:
-                config[key] = {k: v.value for k, v in config[key].items()}

graphdatascience-1.10a1/graphdatascience/graph/graph_remote_proc_runner.py

@@ -0,0 +1,9 @@
+from graphdatascience.graph.base_graph_proc_runner import BaseGraphProcRunner
+from graphdatascience.graph.graph_remote_project_runner import GraphProjectRemoteRunner
+
+
+class GraphRemoteProcRunner(BaseGraphProcRunner):
+    @property
+    def project(self) -> GraphProjectRemoteRunner:
+        self._namespace += ".project.remoteDb"
+        return GraphProjectRemoteRunner(self._query_runner, self._namespace, self._server_version)

graphdatascience-1.10a1/graphdatascience/graph/graph_remote_project_runner.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Any
+
+from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..server_version.compatible_with import compatible_with
+from .graph_object import Graph
+from graphdatascience.call_parameters import CallParameters
+from graphdatascience.graph.graph_create_result import GraphCreateResult
+from graphdatascience.server_version.server_version import ServerVersion
+from graphdatascience.session.schema import (
+    NODE_PROPERTY_SCHEMA,
+    RELATIONSHIP_PROPERTY_SCHEMA,
+)
+
+
+class GraphProjectRemoteRunner(IllegalAttrChecker):
+    @compatible_with("project", min_inclusive=ServerVersion(2, 6, 0))
+    def __call__(self, graph_name: str, query: str, **config: Any) -> GraphCreateResult:
+        placeholder = "<>"  # host and token will be added by query runner
+        self.map_property_types(config)
+        params = CallParameters(
+            graph_name=graph_name,
+            query=query,
+            token=placeholder,
+            host=placeholder,
+            remote_database=self._query_runner.database(),
+            config=config,
+        )
+        result = self._query_runner.call_procedure(
+            endpoint=self._namespace,
+            params=params,
+        ).squeeze()
+        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)
+
+    @staticmethod
+    def map_property_types(config: dict[str, Any]) -> None:
+        for key in [NODE_PROPERTY_SCHEMA, RELATIONSHIP_PROPERTY_SCHEMA]:
+            if key in config:
+                config[key] = {k: v.value for k, v in config[key].items()}

{graphdatascience-1.9 → graphdatascience-1.10a1}/graphdatascience/graph/ogb_loader.py

@@ -314,8 +314,8 @@ class OGBLLoader(OGBLoader):
                 assert source_labels[i] == source_label
                 assert target_labels[i] == target_label
 
-                source_ids[i] += node_id_offsets[edges["head_type"][i]]
-                target_ids[i] += node_id_offsets[edges["tail_type"][i]]
+                source_ids[i] += node_id_offsets[edges["head_type"][i]] + edges["head"][i]
+                target_ids[i] += node_id_offsets[edges["tail_type"][i]] + edges["tail"][i]
 
                 rel_types.append(f"{edge_type}_{set_type.upper()}")
 

graphdatascience-1.10a1/graphdatascience/query_runner/arrow_endpoint_version.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import List
+
+
+class ArrowEndpointVersion(Enum):
+    ALPHA = ""
+    V1 = "v1/"
+
+    def version(self) -> str:
+        return self._name_.lower()
+
+    def prefix(self) -> str:
+        return self._value_
+
+    @staticmethod
+    def from_arrow_info(supported_arrow_versions: List[str]) -> ArrowEndpointVersion:
+        # Fallback for pre 2.6.0 servers that do not support versions
+        if len(supported_arrow_versions) == 0:
+            return ArrowEndpointVersion.ALPHA
+
+        # If the server supports versioned endpoints, we try v1 first
+        if ArrowEndpointVersion.V1.version() in supported_arrow_versions:
+            return ArrowEndpointVersion.V1
+
+        if ArrowEndpointVersion.ALPHA.version() in supported_arrow_versions:
+            return ArrowEndpointVersion.ALPHA
+
+        raise UnsupportedArrowEndpointVersion(supported_arrow_versions)
+
+
+class UnsupportedArrowEndpointVersion(Exception):
+    def __init__(self, server_version: List[str]) -> None:
+        super().__init__(self, f"Unsupported Arrow endpoint versions: {server_version}")

{graphdatascience-1.9 → graphdatascience-1.10a1}/graphdatascience/query_runner/arrow_graph_constructor.py

@@ -11,6 +11,7 @@ from pandas import DataFrame
 from pyarrow import Table
 from tqdm.auto import tqdm
 
+from .arrow_endpoint_version import ArrowEndpointVersion
 from .graph_constructor import GraphConstructor
 
 
@@ -21,6 +22,7 @@ class ArrowGraphConstructor(GraphConstructor):
         graph_name: str,
         flight_client: flight.FlightClient,
         concurrency: int,
+        arrow_endpoint_version: ArrowEndpointVersion,
         undirected_relationship_types: Optional[List[str]],
         chunk_size: int = 10_000,
     ):
@@ -28,6 +30,7 @@ class ArrowGraphConstructor(GraphConstructor):
         self._concurrency = concurrency
         self._graph_name = graph_name
         self._client = flight_client
+        self._arrow_endpoint_version = arrow_endpoint_version
         self._undirected_relationship_types = (
             [] if undirected_relationship_types is None else undirected_relationship_types
         )
@@ -81,6 +84,7 @@ class ArrowGraphConstructor(GraphConstructor):
         return partitioned_dfs
 
     def _send_action(self, action_type: str, meta_data: Dict[str, Any]) -> None:
+        action_type = self._versioned_action_type(action_type)
         result = self._client.do_action(flight.Action(action_type, json.dumps(meta_data).encode("utf-8")))
 
         # Consume result fully to sanity check and avoid cancelled streams
@@ -93,6 +97,7 @@ class ArrowGraphConstructor(GraphConstructor):
         table = Table.from_pandas(df)
         batches = table.to_batches(self._chunk_size)
         flight_descriptor = {"name": self._graph_name, "entity_type": entity_type}
+        flight_descriptor = self._versioned_flight_desriptor(flight_descriptor)
 
         # Write schema
         upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
@@ -117,3 +122,17 @@ class ArrowGraphConstructor(GraphConstructor):
                 if not future.exception():
                     continue
                 raise future.exception()  # type: ignore
+
+    def _versioned_action_type(self, action_type: str) -> str:
+        return self._arrow_endpoint_version.prefix() + action_type
+
+    def _versioned_flight_desriptor(self, flight_descriptor: Dict[str, Any]) -> Dict[str, Any]:
+        return (
+            flight_descriptor
+            if self._arrow_endpoint_version == ArrowEndpointVersion.ALPHA
+            else {
+                "name": "PUT_MESSAGE",
+                "version": ArrowEndpointVersion.V1.version(),
+                "body": flight_descriptor,
+            }
+        )

{graphdatascience-1.9 → graphdatascience-1.10a1}/graphdatascience/query_runner/arrow_query_runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import base64
 import json
 import time
@@ -5,13 +7,14 @@ import warnings
 from typing import Any, Dict, List, Optional, Tuple
 
 import pyarrow.flight as flight
-from pandas import DataFrame, Series
+from pandas import DataFrame
 from pyarrow import ChunkedArray, Table, chunked_array
 from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
 from pyarrow.types import is_dictionary  # type: ignore
 
 from ..call_parameters import CallParameters
 from ..server_version.server_version import ServerVersion
+from .arrow_endpoint_version import ArrowEndpointVersion
 from .arrow_graph_constructor import ArrowGraphConstructor
 from .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
@@ -28,18 +31,14 @@ class ArrowQueryRunner(QueryRunner):
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
-    ) -> "QueryRunner":
+    ) -> QueryRunner:
+        arrow_info = (
+            fallback_query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False).squeeze().to_dict()
+        )
         server_version = fallback_query_runner.server_version()
+        listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])
+        arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info.get("versions", []))
 
-        yield_fields = (
-            ["running", "listenAddress"]
-            if server_version >= ServerVersion(2, 2, 1)
-            else ["running", "advertisedListenAddress"]
-        )
-        arrow_info: "Series[Any]" = fallback_query_runner.call_procedure(
-            endpoint="gds.debug.arrow", yields=yield_fields, custom_error=False
-        ).squeeze()
-        listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])  # type: ignore
         if arrow_info["running"]:
             return ArrowQueryRunner(
                 listen_address,
@@ -49,6 +48,7 @@ class ArrowQueryRunner(QueryRunner):
                 encrypted,
                 disable_server_verification,
                 tls_root_certs,
+                arrow_endpoint_version,
             )
         else:
             return fallback_query_runner
@@ -62,9 +62,11 @@ class ArrowQueryRunner(QueryRunner):
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
+        arrow_endpoint_version: ArrowEndpointVersion = ArrowEndpointVersion.ALPHA,
     ):
         self._fallback_query_runner = fallback_query_runner
         self._server_version = server_version
+        self._arrow_endpoint_version = arrow_endpoint_version
 
         host, port_string = uri.split(":")
 
@@ -272,8 +274,15 @@ class ArrowQueryRunner(QueryRunner):
             "procedure_name": procedure_name,
             "configuration": configuration,
         }
-        ticket = flight.Ticket(json.dumps(payload).encode("utf-8"))
 
+        if self._arrow_endpoint_version == ArrowEndpointVersion.V1:
+            payload = {
+                "name": "GET_MESSAGE",
+                "version": ArrowEndpointVersion.V1.version(),
+                "body": payload,
+            }
+
+        ticket = flight.Ticket(json.dumps(payload).encode("utf-8"))
         get = self._flight_client.do_get(ticket)
         arrow_table = get.read_all()
 
@@ -282,6 +291,13 @@ class ArrowQueryRunner(QueryRunner):
             new_colum_names = ["nodeLabels" if i == "labels" else i for i in arrow_table.column_names]
             arrow_table = arrow_table.rename_columns(new_colum_names)
 
+        # Pandas 2.2.0 deprecated an API used by ArrowTable.to_pandas() (< pyarrow 15.0)
+        warnings.filterwarnings(
+            "ignore",
+            category=DeprecationWarning,
+            message=r"Passing a BlockManager to DataFrame is deprecated",
+        )
+
         return self._sanitize_arrow_table(arrow_table).to_pandas()  # type: ignore
 
     def create_graph_constructor(
@@ -295,7 +311,12 @@ class ArrowQueryRunner(QueryRunner):
             )
 
         return ArrowGraphConstructor(
-            database, graph_name, self._flight_client, concurrency, undirected_relationship_types
+            database,
+            graph_name,
+            self._flight_client,
+            concurrency,
+            self._arrow_endpoint_version,
+            undirected_relationship_types,
         )
 
     def _sanitize_arrow_table(self, arrow_table: Table) -> Table:

{graphdatascience-1.9 → graphdatascience-1.10a1}/graphdatascience/query_runner/aura_db_arrow_query_runner.py

@@ -5,7 +5,7 @@ from pyarrow import flight
 from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
 
 from ..call_parameters import CallParameters
-from ..gds_session.dbms_connection_info import DbmsConnectionInfo
+from ..session.dbms_connection_info import DbmsConnectionInfo
 from .query_runner import QueryRunner
 from graphdatascience.query_runner.graph_constructor import GraphConstructor
 from graphdatascience.server_version.server_version import ServerVersion

graphdatascience-1.10a1/graphdatascience/session/__init__.py

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

{graphdatascience-1.9/graphdatascience/gds_session → graphdatascience-1.10a1/graphdatascience/session}/aura_api.py

@@ -4,8 +4,9 @@ import dataclasses
 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 requests import HTTPError
 from graphdatascience.version import __version__
 
 
-@dataclass(repr=True)
+@dataclass(repr=True, frozen=True)
 class InstanceDetails:
     id: str
     name: str
@@ -31,7 +32,7 @@ class InstanceDetails:
         )
 
 
-@dataclass(repr=True)
+@dataclass(repr=True, frozen=True)
 class InstanceSpecificDetails(InstanceDetails):
     status: str
     connection_url: str
@@ -54,7 +55,7 @@ class InstanceSpecificDetails(InstanceDetails):
         )
 
 
-@dataclass(repr=True)
+@dataclass(repr=True, frozen=True)
 class InstanceCreateDetails:
     id: str
     username: str
@@ -70,6 +71,51 @@ class 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
@@ -87,11 +133,19 @@ 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")
-        self._base_uri = "https://api.neo4j.io" if not self._dev_env else f"https://api-{self._dev_env}.neo4j-dev.io"
+
+        if not self._dev_env:
+            self._base_uri = "https://api.neo4j.io"
+        elif self._dev_env == "staging":
+            self._base_uri = "https://api-staging.neo4j.io"
+        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._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:
@@ -103,14 +157,14 @@ class AuraApi:
         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,
         }
@@ -172,16 +226,16 @@ class AuraApi:
 
     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. "
@@ -191,7 +245,7 @@ class AuraApi:
             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(
@@ -209,6 +263,16 @@ class AuraApi:
 
         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__}"}