graphdatascience 1.11a2__tar.gz → 1.11a3__tar.gz

{graphdatascience-1.11a2/graphdatascience.egg-info → graphdatascience-1.11a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphdatascience
-Version: 1.11a2
+Version: 1.11a3
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -30,8 +30,9 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: multimethod<2.0,>=1.0
 Requires-Dist: neo4j<6.0,>=4.4.2
+Requires-Dist: numpy<2.0
 Requires-Dist: pandas<3.0,>=1.0
-Requires-Dist: pyarrow<16.0,>=11.0
+Requires-Dist: pyarrow<17.0,>=14.0.1
 Requires-Dist: textdistance<5.0,>=4.0
 Requires-Dist: tqdm<5.0,>=4.0
 Requires-Dist: typing-extensions<5.0,>=4.0

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/graph/base_graph_proc_runner.py

@@ -15,12 +15,12 @@ from ..error.uncallable_namespace import UncallableNamespace
 from ..server_version.compatible_with import compatible_with
 from ..server_version.server_version import ServerVersion
 from .graph_entity_ops_runner import (
-    GraphElementPropertyRunner,
     GraphLabelRunner,
     GraphNodePropertiesRunner,
     GraphNodePropertyRunner,
     GraphPropertyRunner,
     GraphRelationshipPropertiesRunner,
+    GraphRelationshipPropertyRunner,
     GraphRelationshipRunner,
     GraphRelationshipsRunner,
 )
@@ -390,9 +390,9 @@ class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
         return GraphNodePropertiesRunner(self._query_runner, self._namespace, self._server_version)
 
     @property
-    def relationshipProperty(self) -> GraphElementPropertyRunner:
+    def relationshipProperty(self) -> GraphRelationshipPropertyRunner:
         self._namespace += ".relationshipProperty"
-        return GraphElementPropertyRunner(self._query_runner, self._namespace, self._server_version)
+        return GraphRelationshipPropertyRunner(self._query_runner, self._namespace, self._server_version)
 
     @property
     def relationshipProperties(self) -> GraphRelationshipPropertiesRunner:

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/graph/graph_entity_ops_runner.py

@@ -70,13 +70,6 @@ class GraphEntityOpsBaseRunner(UncallableNamespace, IllegalAttrChecker):
         )
 
 
-class GraphElementPropertyRunner(GraphEntityOpsBaseRunner):
-    @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
-    def stream(self, G: Graph, node_properties: str, node_labels: Strings = ["*"], **config: Any) -> DataFrame:
-        self._namespace += ".stream"
-        return self._handle_properties(G, node_properties, node_labels, config)
-
-
 class GraphNodePropertyRunner(GraphEntityOpsBaseRunner):
     @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
     @filter_id_func_deprecation_warning()
@@ -197,6 +190,16 @@ class GraphNodePropertiesRunner(GraphEntityOpsBaseRunner):
         ).squeeze()
 
 
+class GraphRelationshipPropertyRunner(GraphEntityOpsBaseRunner):
+    @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
+    def stream(
+        self, G: Graph, relationship_property: str, relationship_types: Strings = ["*"], **config: Any
+    ) -> DataFrame:
+        self._namespace += ".stream"
+        relationship_types = [relationship_types] if isinstance(relationship_types, str) else relationship_types
+        return self._handle_properties(G, relationship_property, relationship_types, config)
+
+
 class GraphRelationshipPropertiesRunner(GraphEntityOpsBaseRunner):
     @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
     def stream(
@@ -209,6 +212,8 @@ class GraphRelationshipPropertiesRunner(GraphEntityOpsBaseRunner):
     ) -> DataFrame:
         self._namespace += ".stream"
 
+        relationship_types = [relationship_types] if isinstance(relationship_types, str) else relationship_types
+
         result = self._handle_properties(G, relationship_properties, relationship_types, config)
 
         # new format was requested, but the query was run via Cypher

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/graph/graph_remote_project_runner.py

@@ -20,18 +20,24 @@ class GraphProjectRemoteRunner(IllegalAttrChecker):
         concurrency: int = 4,
         undirected_relationship_types: Optional[List[str]] = None,
         inverse_indexed_relationship_types: Optional[List[str]] = None,
+        batch_size: Optional[int] = None,
     ) -> GraphCreateResult:
         if inverse_indexed_relationship_types is None:
             inverse_indexed_relationship_types = []
         if undirected_relationship_types is None:
             undirected_relationship_types = []
 
+        arrow_configuration = {}
+        if batch_size is not None:
+            arrow_configuration["batchSize"] = batch_size
+
         params = CallParameters(
             graph_name=graph_name,
             query=query,
             concurrency=concurrency,
             undirected_relationship_types=undirected_relationship_types,
             inverse_indexed_relationship_types=inverse_indexed_relationship_types,
+            arrow_configuration=arrow_configuration,
         )
 
         result = self._query_runner.call_procedure(

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/query_runner/aura_db_query_runner.py

@@ -1,5 +1,6 @@
 import time
 from typing import Any, Dict, List, Optional
+from uuid import uuid4
 
 from pandas import DataFrame
 
@@ -47,7 +48,7 @@ class AuraDbQueryRunner(QueryRunner):
             params = CallParameters()
 
         if AuraDbQueryRunner.GDS_REMOTE_PROJECTION_PROC_NAME == endpoint:
-            return self._remote_projection(endpoint, params, yields, database, logging, custom_error)
+            return self._remote_projection(endpoint, params, yields, database, logging)
 
         elif ".write" in endpoint and self.is_remote_projected_graph(params["graph_name"]):
             return self._remote_write_back(endpoint, params, yields, database, logging, custom_error)
@@ -103,9 +104,8 @@ class AuraDbQueryRunner(QueryRunner):
         yields: Optional[List[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
-        custom_error: bool = True,
     ) -> DataFrame:
-        self._inject_connection_parameters(params)
+        self._inject_arrow_config(params["arrow_configuration"])
         return self._db_query_runner.call_procedure(endpoint, params, yields, database, logging, False)
 
     def _remote_write_back(
@@ -120,21 +120,29 @@ class AuraDbQueryRunner(QueryRunner):
         if params["config"] is None:
             params["config"] = {}
 
+        # we pop these out so that they are not retained for the GDS proc call
+        db_arrow_config = params["config"].pop("arrowConfiguration", {})  # type: ignore
+        self._inject_arrow_config(db_arrow_config)
+
+        job_id = params["config"]["jobId"] if "jobId" in params["config"] else str(uuid4())  # type: ignore
+        params["config"]["jobId"] = job_id  # type: ignore
+
         params["config"]["writeToResultStore"] = True  # type: ignore
+
         gds_write_result = self._gds_query_runner.call_procedure(
             endpoint, params, yields, database, logging, custom_error
         )
 
-        write_params = {
+        db_write_proc_params = {
             "graphName": params["graph_name"],
             "databaseName": self._gds_query_runner.database(),
-            "writeConfiguration": self._extract_write_back_arguments(endpoint, params),
+            "jobId": job_id,
+            "arrowConfiguration": db_arrow_config,
         }
-        self._inject_connection_parameters(write_params)
 
         write_back_start = time.time()
         database_write_result = self._db_query_runner.call_procedure(
-            "gds.arrow.write", CallParameters(write_params), yields, None, False, False
+            "gds.arrow.write", CallParameters(db_write_proc_params), yields, None, False, False
         )
         write_millis = (time.time() - write_back_start) * 1000
         gds_write_result["writeMillis"] = write_millis
@@ -150,75 +158,13 @@ class AuraDbQueryRunner(QueryRunner):
 
         return gds_write_result
 
-    def _inject_connection_parameters(self, params: Dict[str, Any]) -> None:
+    def _inject_arrow_config(self, params: Dict[str, Any]) -> None:
         host, port = self._gds_arrow_client.connection_info()
         token = self._gds_arrow_client.request_token()
         if token is None:
             token = "IGNORED"
-        params["arrowConfiguration"] = {
-            "host": host,
-            "port": port,
-            "token": token,
-            "encrypted": self._encrypted,
-        }
 
-    @staticmethod
-    def _extract_write_back_arguments(proc_name: str, params: Dict[str, Any]) -> Dict[str, Any]:
-        config = params.get("config", {})
-        write_config = {}
-
-        if "writeConcurrency" in config:
-            write_config["concurrency"] = config["writeConcurrency"]
-        elif "concurrency" in config:
-            write_config["concurrency"] = config["concurrency"]
-
-        if "gds.shortestPath" in proc_name or "gds.allShortestPaths" in proc_name:
-            write_config["relationshipType"] = config["writeRelationshipType"]
-
-            write_node_ids = config.get("writeNodeIds")
-            write_costs = config.get("writeCosts")
-
-            if write_node_ids and write_costs:
-                write_config["relationshipProperties"] = ["totalCost", "nodeIds", "costs"]
-            elif write_node_ids:
-                write_config["relationshipProperties"] = ["totalCost", "nodeIds"]
-            elif write_costs:
-                write_config["relationshipProperties"] = ["totalCost", "costs"]
-            else:
-                write_config["relationshipProperties"] = ["totalCost"]
-
-        elif "gds.graph." in proc_name:
-            if "gds.graph.nodeProperties.write" == proc_name:
-                properties = params["properties"]
-                write_config["nodeProperties"] = properties if isinstance(properties, list) else [properties]
-                write_config["nodeLabels"] = params["entities"]
-
-            elif "gds.graph.nodeLabel.write" == proc_name:
-                write_config["nodeLabels"] = [params["node_label"]]
-
-            elif "gds.graph.relationshipProperties.write" == proc_name:
-                write_config["relationshipProperties"] = params["relationship_properties"]
-                write_config["relationshipType"] = params["relationship_type"]
-
-            elif "gds.graph.relationship.write" == proc_name:
-                if "relationship_property" in params and params["relationship_property"] != "":
-                    write_config["relationshipProperties"] = [params["relationship_property"]]
-                write_config["relationshipType"] = params["relationship_type"]
-
-            else:
-                raise ValueError(f"Unsupported procedure name: {proc_name}")
-
-        else:
-            if "writeRelationshipType" in config:
-                write_config["relationshipType"] = config["writeRelationshipType"]
-                if "writeProperty" in config:
-                    write_config["relationshipProperties"] = [config["writeProperty"]]
-            else:
-                if "writeProperty" in config:
-                    write_config["nodeProperties"] = [config["writeProperty"]]
-                if "nodeLabels" in params:
-                    write_config["nodeLabels"] = params["nodeLabels"]
-                else:
-                    write_config["nodeLabels"] = ["*"]
-
-        return write_config
+        params["host"] = host
+        params["port"] = port
+        params["token"] = token
+        params["encrypted"] = self._encrypted

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/query_runner/gds_arrow_client.py

@@ -18,8 +18,10 @@ from .query_runner import QueryRunner
 class GdsArrowClient:
     @staticmethod
     def is_arrow_enabled(query_runner: QueryRunner) -> bool:
-        arrow_info = query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False).squeeze().to_dict()
-        return not not arrow_info["running"]
+        arrow_server_running = query_runner.call_procedure(
+            endpoint="gds.debug.arrow", custom_error=False, yields=["running", "listenAddress"]
+        ).squeeze()
+        return not not arrow_server_running["running"]
 
     @staticmethod
     def create(
@@ -30,14 +32,24 @@ class GdsArrowClient:
         tls_root_certs: Optional[bytes] = None,
         connection_string_override: Optional[str] = None,
     ) -> "GdsArrowClient":
-        arrow_info = query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False).squeeze().to_dict()
+        if query_runner.server_version() > ServerVersion(2, 6, 0):
+            debugYields = ["listenAddress", "versions"]
+        else:
+            # return enabled to avoid squeeze returning a single str
+            debugYields = ["listenAddress", "enabled"]
+
+        arrow_info = (
+            query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False, yields=debugYields)
+            .squeeze()
+            .to_dict()
+        )
 
         server_version = query_runner.server_version()
         connection_string: str
         if connection_string_override is not None:
             connection_string = connection_string_override
         else:
-            connection_string = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])
+            connection_string = arrow_info.get("listenAddress")
 
         host, port = connection_string.split(":")
 

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/session/__init__.py

@@ -1,9 +1,8 @@
 from .algorithm_category import AlgorithmCategory
 from .dbms_connection_info import DbmsConnectionInfo
 from .gds_sessions import AuraAPICredentials, GdsSessions
-from .schema import GdsPropertyTypes
 from .session_info import SessionInfo
-from .session_sizes import SessionMemory
+from .session_sizes import SessionMemory, SessionMemoryValue
 
 __all__ = [
     "GdsSessions",
@@ -11,6 +10,6 @@ __all__ = [
     "DbmsConnectionInfo",
     "AuraAPICredentials",
     "SessionMemory",
-    "GdsPropertyTypes",
+    "SessionMemoryValue",
     "AlgorithmCategory",
 ]

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/session/aura_api.py

@@ -19,6 +19,7 @@ from graphdatascience.session.aura_api_responses import (
     TenantDetails,
     WaitResult,
 )
+from graphdatascience.session.session_sizes import SessionMemoryValue
 from graphdatascience.version import __version__
 
 
@@ -62,11 +63,11 @@ class AuraApi:
 
         return host.split(".")[0].split("-")[0]
 
-    def create_session(self, name: str, dbid: str, pwd: str, memory: str) -> SessionDetails:
+    def create_session(self, name: str, dbid: str, pwd: str, memory: SessionMemoryValue) -> SessionDetails:
         response = req.post(
             f"{self._base_uri}/v1beta5/data-science/sessions",
             headers=self._build_header(),
-            json={"name": name, "instance_id": dbid, "password": pwd, "memory": memory},
+            json={"name": name, "instance_id": dbid, "password": pwd, "memory": memory.value},
         )
 
         response.raise_for_status()
@@ -141,12 +142,14 @@ class AuraApi:
 
         return False
 
-    def create_instance(self, name: str, memory: str, cloud_provider: str, region: str) -> InstanceCreateDetails:
+    def create_instance(
+        self, name: str, memory: SessionMemoryValue, cloud_provider: str, region: str
+    ) -> InstanceCreateDetails:
         tenant_details = self.tenant_details()
 
         data = {
             "name": name,
-            "memory": memory,
+            "memory": memory.value,
             "version": "5",
             "region": region,
             "type": tenant_details.ds_type,

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/session/aura_api_responses.py

@@ -9,13 +9,15 @@ from typing import Any, Dict, NamedTuple, Optional, Set
 
 from pandas import Timedelta
 
+from .session_sizes import SessionMemoryValue
+
 
 @dataclass(repr=True, frozen=True)
 class SessionDetails:
     id: str
     name: str
     instance_id: str
-    memory: str
+    memory: SessionMemoryValue
     status: str
     host: str
     created_at: datetime
@@ -31,7 +33,7 @@ class SessionDetails:
             id=json["id"],
             name=json["name"],
             instance_id=json["instance_id"],
-            memory=json["memory"],
+            memory=SessionMemoryValue.fromApiResponse(json["memory"]),
             status=json["status"],
             host=json["host"],
             expiry_date=TimeParser.fromisoformat(expiry_date) if expiry_date else None,
@@ -40,7 +42,7 @@ class SessionDetails:
         )
 
     def bolt_connection_url(self) -> str:
-        return f"neo4j+ssc://{self.host}"  # TODO use neo4j+s
+        return f"neo4j+s://{self.host}"
 
     def is_expired(self) -> bool:
         return self.status == "Expired"
@@ -67,7 +69,7 @@ class InstanceDetails:
 class InstanceSpecificDetails(InstanceDetails):
     status: str
     connection_url: str
-    memory: str
+    memory: SessionMemoryValue
     type: str
     region: str
 
@@ -80,7 +82,7 @@ class InstanceSpecificDetails(InstanceDetails):
             cloud_provider=json["cloud_provider"],
             status=json["status"],
             connection_url=json.get("connection_url", ""),
-            memory=json.get("memory", ""),
+            memory=SessionMemoryValue.fromInstanceSize(json.get("memory")),
             type=json["type"],
             region=json["region"],
         )

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/session/aurads_sessions.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import warnings
 from typing import List, Optional
 
@@ -16,7 +17,11 @@ from graphdatascience.session.aura_graph_data_science import AuraGraphDataScienc
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
 from graphdatascience.session.region_suggester import closest_match
 from graphdatascience.session.session_info import SessionInfo
-from graphdatascience.session.session_sizes import SessionMemory
+from graphdatascience.session.session_sizes import (
+    SESSION_MEMORY_VALUE_UNKNOWN,
+    SessionMemory,
+    SessionMemoryValue,
+)
 
 
 class AuraDsSessions:
@@ -25,6 +30,7 @@ class AuraDsSessions:
 
     def __init__(self, aura_api: AuraApi) -> None:
         self._aura_api = aura_api
+        self._logger = logging.getLogger()
 
     def estimate(
         self, node_count: int, relationship_count: int, algorithm_categories: Optional[List[AlgorithmCategory]] = None
@@ -41,7 +47,7 @@ class AuraDsSessions:
                 ResourceWarning,
             )
 
-        return SessionMemory(estimation.recommended_size)
+        return SessionMemory(SessionMemoryValue(estimation.recommended_size))
 
     def get_or_create(
         self,
@@ -53,14 +59,18 @@ class AuraDsSessions:
 
         if existing_session:
             session_id = existing_session.id
-            # 0MB is AuraAPI default value for memory if none can be retrieved
-            if existing_session.memory != "0MB" and existing_session.memory != memory.value:
+            # AuraAPI default value for memory if none can be retrieved
+            if existing_session.memory == SESSION_MEMORY_VALUE_UNKNOWN:
+                self._logger.debug(
+                    f"Reusing existing session with id `{session_id}` as size for session is unknown during creation."
+                )
+            elif existing_session.memory != memory.value:
                 raise ValueError(
-                    f"Session `{session_name}` already exists with memory `{existing_session.memory}`. "
+                    f"Session `{session_name}` already exists with memory `{existing_session.memory.value}`. "
                     f"Requested memory `{memory.value}` does not match."
                 )
         else:
-            create_details = self._create_session(session_name, memory, db_connection)
+            create_details = self._create_session(session_name, memory.value, db_connection)
             session_id = create_details.id
 
         wait_result = self._aura_api.wait_for_instance_running(session_id)
@@ -118,7 +128,7 @@ class AuraDsSessions:
         return self._aura_api.list_instance(matched_instances[0].id)
 
     def _create_session(
-        self, session_name: str, memory: SessionMemory, db_connection: DbmsConnectionInfo
+        self, session_name: str, memory: SessionMemoryValue, db_connection: DbmsConnectionInfo
     ) -> InstanceCreateDetails:
         db_instance_id = AuraApi.extract_id(db_connection.uri)
         db_instance = self._aura_api.list_instance(db_instance_id)
@@ -128,7 +138,7 @@ class AuraDsSessions:
         region = self._ds_region(db_instance.region, db_instance.cloud_provider)
 
         create_details = self._aura_api.create_instance(
-            SessionNameHelper.instance_name(session_name), memory.value, db_instance.cloud_provider, region
+            SessionNameHelper.instance_name(session_name), memory, db_instance.cloud_provider, region
         )
         return create_details
 

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/session/dedicated_sessions.py

@@ -5,13 +5,15 @@ import warnings
 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_responses import SessionDetails
 from graphdatascience.session.aura_graph_data_science import AuraGraphDataScience
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
 from graphdatascience.session.session_info import SessionInfo
-from graphdatascience.session.session_sizes import SessionMemory
+from graphdatascience.session.session_sizes import SessionMemory, SessionMemoryValue
 
 
 class DedicatedSessions:
@@ -35,7 +37,7 @@ class DedicatedSessions:
                 ResourceWarning,
             )
 
-        return SessionMemory(estimation.recommended_size)
+        return SessionMemory(SessionMemoryValue(estimation.recommended_size))
 
     def get_or_create(
         self,
@@ -52,9 +54,10 @@ class DedicatedSessions:
         # TODO configure session size (and check existing_session has same size)
         if existing_session:
             self._check_expiry_date(existing_session)
+            self._check_memory_configuration(existing_session, memory.value)
             session_id = existing_session.id
         else:
-            create_details = self._create_session(session_name, dbid, db_connection.uri, password, memory)
+            create_details = self._create_session(session_name, dbid, db_connection.uri, password, memory.value)
             session_id = create_details.id
 
         wait_result = self._aura_api.wait_for_session_running(session_id, dbid)
@@ -78,12 +81,13 @@ class DedicatedSessions:
             for db in dbs:
                 candidate = self._find_existing_session(session_name, db.id)
                 if candidate:
+                    dbid = db.id
                     break
         else:
             candidate = self._find_existing_session(session_name, dbid)
 
         if candidate:
-            return self._aura_api.delete_session(candidate.id, db.id) is not None
+            return self._aura_api.delete_session(candidate.id, str(dbid)) is not None
 
         return False
 
@@ -92,7 +96,12 @@ class DedicatedSessions:
 
         sessions: List[SessionDetails] = []
         for db in dbs:
-            sessions.extend(self._aura_api.list_sessions(db.id))
+            try:
+                sessions.extend(self._aura_api.list_sessions(db.id))
+            except HTTPError as e:
+                # ignore 404 errors when listing sessions as it could mean paused sessions or deleted sessions
+                if e.response.status_code != 404:
+                    raise e
 
         return [SessionInfo.from_session_details(i) for i in sessions]
 
@@ -108,7 +117,7 @@ class DedicatedSessions:
         return matched_sessions[0]
 
     def _create_session(
-        self, session_name: str, dbid: str, dburi: str, pwd: str, memory: SessionMemory
+        self, session_name: str, dbid: str, dburi: str, pwd: str, memory: SessionMemoryValue
     ) -> SessionDetails:
         db_instance = self._aura_api.list_instance(dbid)
         if not db_instance:
@@ -118,7 +127,7 @@ class DedicatedSessions:
             name=session_name,
             dbid=dbid,
             pwd=pwd,
-            memory=memory.value,
+            memory=memory,
         )
         return create_details
 
@@ -139,6 +148,15 @@ class DedicatedSessions:
             if until_expiry < timedelta(days=1):
                 raise Warning(f"Session `{session.name}` is expiring in less than a day.")
 
+    def _check_memory_configuration(
+        self, existing_session: SessionDetails, requested_memory: SessionMemoryValue
+    ) -> None:
+        if existing_session.memory != requested_memory:
+            raise RuntimeError(
+                f"Session `{existing_session.name}` exists with a different memory configuration. "
+                f"Current: {existing_session.memory}, Requested: {requested_memory}."
+            )
+
     @classmethod
     def _fail_ambiguous_session(cls, session_name: str, sessions: List[SessionDetails]) -> None:
         candidates = [i.id for i in sessions]

{graphdatascience-1.11a2 → graphdatascience-1.11a3}/graphdatascience/session/session_info.py

@@ -5,6 +5,7 @@ from datetime import datetime
 from typing import Optional
 
 from graphdatascience.session.aura_api_responses import SessionDetails
+from graphdatascience.session.session_sizes import SessionMemoryValue
 
 
 @dataclass(frozen=True)
@@ -18,7 +19,7 @@ class SessionInfo:
     """
 
     name: str
-    memory: str
+    memory: SessionMemoryValue
 
     @classmethod
     def from_session_details(cls, details: SessionDetails) -> ExtendedSessionInfo:

graphdatascience-1.11a3/graphdatascience/session/session_sizes.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Optional
+
+
+@dataclass(frozen=True)
+class SessionMemoryValue:
+    value: str
+
+    def __str__(self) -> str:
+        return self.value
+
+    @staticmethod
+    def fromApiResponse(value: str) -> SessionMemoryValue:
+        """
+        Converts the string value from an API response to a SessionMemory enumeration value.
+
+        Args:
+            value: The string value from the API response.
+
+        Returns:
+            The SessionMemory enumeration value.
+
+        """
+        if value == "":
+            raise ValueError("memory configuration cannot be empty")
+
+        return SessionMemoryValue(value.replace("Gi", "GB"))
+
+    @staticmethod
+    def fromInstanceSize(value: Optional[str]) -> SessionMemoryValue:
+        if not value:
+            return SESSION_MEMORY_VALUE_UNKNOWN
+
+        return SessionMemoryValue(value.replace("Gi", "GB"))
+
+
+SESSION_MEMORY_VALUE_UNKNOWN = SessionMemoryValue("")
+
+
+class SessionMemory(Enum):
+    """
+    Enumeration representing session main memory configurations.
+    """
+
+    m_4GB = SessionMemoryValue("4GB")
+    m_8GB = SessionMemoryValue("8GB")
+    m_16GB = SessionMemoryValue("16GB")
+    m_24GB = SessionMemoryValue("24GB")
+    m_32GB = SessionMemoryValue("32GB")
+    m_48GB = SessionMemoryValue("48GB")
+    m_64GB = SessionMemoryValue("64GB")
+    m_96GB = SessionMemoryValue("96GB")
+    m_128GB = SessionMemoryValue("128GB")
+    m_192GB = SessionMemoryValue("192GB")
+    m_256GB = SessionMemoryValue("256GB")
+    m_384GB = SessionMemoryValue("384GB")
+
+    @classmethod
+    def all_values(cls) -> List[SessionMemoryValue]:
+        """
+        All supported memory configurations.
+
+        Returns:
+            A list of strings representing all supported memory configurations for sessions.
+
+        """
+        return [e.value for e in cls]

graphdatascience-1.11a3/graphdatascience/version.py

@@ -0,0 +1 @@
+__version__ = "1.11a3"