Use a data class instead of an enum to represent session memory

* We will keep the enum to make the selection easier for users

Author: DarthMax

graphdatascience/session/aura_api.py

@@ -9,7 +9,6 @@
 import requests as req
 from requests import HTTPError
 
-from graphdatascience.session.session_sizes import SessionMemory
 from graphdatascience.session.algorithm_category import AlgorithmCategory
 from graphdatascience.session.aura_api_responses import (
     EstimationDetails,
@@ -20,6 +19,7 @@
     TenantDetails,
     WaitResult,
 )
+from graphdatascience.session.session_sizes import SessionMemoryValue
 from graphdatascience.version import __version__
 
 
@@ -63,7 +63,7 @@ def extract_id(uri: str) -> str:
 
         return host.split(".")[0].split("-")[0]
 
-    def create_session(self, name: str, dbid: str, pwd: str, memory: SessionMemory) -> 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(),
@@ -143,7 +143,7 @@ def delete_session(self, session_id: str, dbid: str) -> bool:
         return False
 
     def create_instance(
-        self, name: str, memory: SessionMemory, cloud_provider: str, region: str
+        self, name: str, memory: SessionMemoryValue, cloud_provider: str, region: str
     ) -> InstanceCreateDetails:
         tenant_details = self.tenant_details()
 

graphdatascience/session/aura_api_responses.py

@@ -9,15 +9,15 @@
 
 from pandas import Timedelta
 
-from .session_sizes import SessionMemory
+from .session_sizes import SessionMemoryValue
 
 
 @dataclass(repr=True, frozen=True)
 class SessionDetails:
     id: str
     name: str
     instance_id: str
-    memory: SessionMemory
+    memory: SessionMemoryValue
     status: str
     host: str
     created_at: datetime
@@ -33,7 +33,7 @@ def fromJson(cls, json: Dict[str, Any]) -> SessionDetails:
             id=json["id"],
             name=json["name"],
             instance_id=json["instance_id"],
-            memory=SessionMemory.fromApiResponse(json["memory"]),
+            memory=SessionMemoryValue.fromApiResponse(json["memory"]),
             status=json["status"],
             host=json["host"],
             expiry_date=TimeParser.fromisoformat(expiry_date) if expiry_date else None,
@@ -69,7 +69,7 @@ def fromJson(cls, json: Dict[str, Any]) -> InstanceDetails:
 class InstanceSpecificDetails(InstanceDetails):
     status: str
     connection_url: str
-    memory: SessionMemory
+    memory: SessionMemoryValue
     type: str
     region: str
 
@@ -82,7 +82,7 @@ def fromJson(cls, json: Dict[str, Any]) -> InstanceSpecificDetails:
             cloud_provider=json["cloud_provider"],
             status=json["status"],
             connection_url=json.get("connection_url", ""),
-            memory=SessionMemory.fromApiResponse(json.get("memory", "")),
+            memory=SessionMemoryValue.fromApiResponse(json.get("memory", "")),
             type=json["type"],
             region=json["region"],
         )

graphdatascience/session/aurads_sessions.py

@@ -16,7 +16,7 @@
 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 SessionMemory, SessionMemoryValue
 
 
 class AuraDsSessions:
@@ -41,7 +41,7 @@ def estimate(
                 ResourceWarning,
             )
 
-        return SessionMemory(estimation.recommended_size)
+        return SessionMemory(SessionMemoryValue(estimation.recommended_size))
 
     def get_or_create(
         self,
@@ -54,13 +54,13 @@ def get_or_create(
         if existing_session:
             session_id = existing_session.id
             # 0MB is AuraAPI default value for memory if none can be retrieved
-            if existing_session.memory != memory:
+            if existing_session.memory.value != "0MB" and existing_session.memory != memory.value:
                 raise ValueError(
                     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 +118,7 @@ def _find_existing_session(self, session_name: str) -> Optional[InstanceSpecific
         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)

graphdatascience/session/dedicated_sessions.py

@@ -11,7 +11,7 @@
 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 +35,7 @@ def estimate(
                 ResourceWarning,
             )
 
-        return SessionMemory(estimation.recommended_size)
+        return SessionMemory(SessionMemoryValue(estimation.recommended_size))
 
     def get_or_create(
         self,
@@ -52,10 +52,10 @@ def get_or_create(
         # 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)
+            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)
@@ -109,7 +109,7 @@ def _find_existing_session(self, session_name: str, dbid: str) -> Optional[Sessi
         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:
@@ -140,11 +140,13 @@ def _check_expiry_date(self, session: SessionDetails) -> None:
             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: SessionMemory) -> None:
+    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.value}, Requested: {requested_memory.value}."
+                f"Current: {existing_session.memory}, Requested: {requested_memory}."
             )
 
     @classmethod

graphdatascience/session/session_info.py

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

graphdatascience/session/session_sizes.py

@@ -1,27 +1,53 @@
+from dataclasses import dataclass
 from enum import Enum
 from typing import List
 
 
+@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"))
+
+
 class SessionMemory(Enum):
     """
     Enumeration representing session main memory configurations.
     """
 
-    m_4GB = "4GB"
-    m_8GB = "8GB"
-    m_16GB = "16GB"
-    m_24GB = "24GB"
-    m_32GB = "32GB"
-    m_48GB = "48GB"
-    m_64GB = "64GB"
-    m_96GB = "96GB"
-    m_128GB = "128GB"
-    m_192GB = "192GB"
-    m_256GB = "256GB"
-    m_384GB = "384GB"
+    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[str]:
+    def all_values(cls) -> List[SessionMemoryValue]:
         """
         All supported memory configurations.
 
@@ -30,20 +56,3 @@ def all_values(cls) -> List[str]:
 
         """
         return [e.value for e in cls]
-
-    @staticmethod
-    def fromApiResponse(value: str) -> "SessionMemory":
-        """
-        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.
-
-        """
-        try:
-            return SessionMemory(value.replace("Gi", "GB"))
-        except ValueError:
-            raise ValueError(f"Unsupported memory configuration: {value}")

graphdatascience/tests/unit/test_aura_api.py

@@ -38,7 +38,7 @@ def test_create_session(requests_mock: Mocker) -> None:
         },
     )
 
-    result = api.create_session("name-0", "dbid-1", "pwd-2", SessionMemory.m_4GB)
+    result = api.create_session("name-0", "dbid-1", "pwd-2", SessionMemory.m_4GB.value)
 
     assert result == SessionDetails(
         id="id0",
@@ -47,7 +47,7 @@ def test_create_session(requests_mock: Mocker) -> None:
         instance_id="dbid-1",
         created_at=TimeParser.fromisoformat("1970-01-01T00:00:00Z"),
         host="1.2.3.4",
-        memory=SessionMemory.m_4GB,
+        memory=SessionMemory.m_4GB.value,
         expiry_date=None,
         ttl=None,
     )
@@ -82,7 +82,7 @@ def test_list_session(requests_mock: Mocker) -> None:
         instance_id="dbid-1",
         created_at=TimeParser.fromisoformat("1970-01-01T00:00:00Z"),
         host="1.2.3.4",
-        memory=SessionMemory.m_4GB,
+        memory=SessionMemory.m_4GB.value,
         expiry_date=TimeParser.fromisoformat("1977-01-01T00:00:00Z"),
         ttl=None,
     )
@@ -126,7 +126,7 @@ def test_list_sessions(requests_mock: Mocker) -> None:
         instance_id="dbid-1",
         created_at=TimeParser.fromisoformat("1970-01-01T00:00:00Z"),
         host="1.2.3.4",
-        memory=SessionMemory.m_4GB,
+        memory=SessionMemory.m_4GB.value,
         expiry_date=TimeParser.fromisoformat("1977-01-01T00:00:00Z"),
         ttl=None,
     )
@@ -137,7 +137,7 @@ def test_list_sessions(requests_mock: Mocker) -> None:
         status="Creating",
         instance_id="dbid-3",
         created_at=TimeParser.fromisoformat("2012-01-01T00:00:00Z"),
-        memory=SessionMemory.m_8GB,
+        memory=SessionMemory.m_8GB.value,
         host="foo.bar",
         expiry_date=None,
         ttl=None,
@@ -270,7 +270,7 @@ def test_delete_instance(requests_mock: Mocker) -> None:
         cloud_provider="",
         status="deleting",
         connection_url="",
-        memory=SessionMemory.m_4GB,
+        memory=SessionMemory.m_4GB.value,
         region="",
         type="",
     )
@@ -337,7 +337,7 @@ def test_create_instance(requests_mock: Mocker) -> None:
         },
     )
 
-    api.create_instance("name", SessionMemory.m_16GB, "gcp", "leipzig-1")
+    api.create_instance("name", SessionMemory.m_16GB.value, "gcp", "leipzig-1")
 
     requested_data = requests_mock.request_history[-1].json()
     assert requested_data["name"] == "name"
@@ -446,7 +446,7 @@ def test_list_instance_missing_memory_field(requests_mock: Mocker) -> None:
     result = api.list_instance("id0")
 
     assert result and result.id == "a10fb995"
-    assert result.memory == SessionMemory.m_16GB
+    assert result.memory == SessionMemory.m_16GB.value
 
 
 def test_list_missing_instance(requests_mock: Mocker) -> None:
@@ -653,7 +653,7 @@ def test_parse_session_info() -> None:
     assert session_info == SessionDetails(
         id="test_id",
         name="test_session",
-        memory=SessionMemory.m_4GB,
+        memory=SessionMemory.m_4GB.value,
         instance_id="test_instance",
         status="running",
         host="a.b",
@@ -678,7 +678,7 @@ def test_parse_session_info_without_optionals() -> None:
     assert session_info == SessionDetails(
         id="test_id",
         name="test_session",
-        memory=SessionMemory.m_16GB,
+        memory=SessionMemory.m_16GB.value,
         instance_id="test_instance",
         host="a.b",
         status="running",