Merge pull request #643 from DarthMax/support_session_sizing

Forward memory for dedicated sessions

graphdatascience/session/aura_api.py

@@ -62,11 +62,11 @@ def extract_id(uri: str) -> str:
 
         return host.split(".")[0].split("-")[0]
 
-    def create_session(self, name: str, dbid: str, pwd: str) -> SessionDetails:
+    def create_session(self, name: str, dbid: str, pwd: str, memory: str) -> SessionDetails:
         response = req.post(
             f"{self._base_uri}/v1beta5/data-science/sessions",
             headers=self._build_header(),
-            json={"name": name, "instance_id": dbid, "password": pwd},
+            json={"name": name, "instance_id": dbid, "password": pwd, "memory": memory},
         )
 
         response.raise_for_status()

graphdatascience/session/dedicated_sessions.py

@@ -52,7 +52,7 @@ def get_or_create(
         if existing_session:
             session_id = existing_session.id
         else:
-            create_details = self._create_session(session_name, dbid, db_connection.uri, password)
+            create_details = self._create_session(session_name, dbid, db_connection.uri, password, memory)
             session_id = create_details.id
 
         wait_result = self._aura_api.wait_for_session_running(session_id, dbid)
@@ -105,7 +105,9 @@ 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) -> SessionDetails:
+    def _create_session(
+        self, session_name: str, dbid: str, dburi: str, pwd: str, memory: SessionMemory
+    ) -> SessionDetails:
         db_instance = self._aura_api.list_instance(dbid)
         if not db_instance:
             raise ValueError(f"Could not find AuraDB instance with the uri `{dburi}`")
@@ -114,6 +116,7 @@ def _create_session(self, session_name: str, dbid: str, dburi: str, pwd: str) ->
             name=session_name,
             dbid=dbid,
             pwd=pwd,
+            memory=memory.value,
         )
         return create_details
 

graphdatascience/tests/unit/test_aura_api.py

@@ -37,7 +37,7 @@ def test_create_session(requests_mock: Mocker) -> None:
         },
     )
 
-    result = api.create_session("name-0", "dbid-1", "pwd-2")
+    result = api.create_session("name-0", "dbid-1", "pwd-2", "4G")
 
     assert result == SessionDetails(
         id="id0",

graphdatascience/tests/unit/test_dedicated_sessions.py

@@ -43,12 +43,12 @@ def __init__(
         self._status_after_creating = status_after_creating
         self._size_estimation = size_estimation or EstimationDetails("1GB", "8GB", False)
 
-    def create_session(self, name: str, dbid: str, pwd: str) -> SessionDetails:
+    def create_session(self, name: str, dbid: str, pwd: str, memory: str) -> SessionDetails:
         details = SessionDetails(
             id=f"{dbid}-ffff{self.id_counter}",
             name=name,
             instance_id=dbid,
-            memory="2GB",
+            memory=memory,
             status="Creating",
             created_at=datetime.fromisoformat("2021-01-01T00:00:00+00:00"),
             host="foo.bar",
@@ -148,7 +148,7 @@ def aura_api() -> AuraApi:
 def test_list_session(aura_api: AuraApi) -> None:
     _setup_db_instance(aura_api)
     session = aura_api.create_session(
-        name="gds-session-my-session-name", dbid=aura_api.list_instances()[0].id, pwd="some_pwd"
+        name="gds-session-my-session-name", dbid=aura_api.list_instances()[0].id, pwd="some_pwd", memory="8GB"
     )
     sessions = DedicatedSessions(aura_api)
 
@@ -214,8 +214,8 @@ def test_get_or_create(mocker: MockerFixture, aura_api: AuraApi) -> None:
 
 def test_get_or_create_duplicate_session(aura_api: AuraApi) -> None:
     db = _setup_db_instance(aura_api)
-    aura_api.create_session("one", db.id, "1234")
-    aura_api.create_session("one", db.id, "12345")
+    aura_api.create_session("one", db.id, "1234", memory="1GB")
+    aura_api.create_session("one", db.id, "12345", memory="1GB")
 
     sessions = DedicatedSessions(aura_api)
 
@@ -226,8 +226,8 @@ def test_get_or_create_duplicate_session(aura_api: AuraApi) -> None:
 def test_delete_session(aura_api: AuraApi) -> None:
     db1 = aura_api.create_instance("db1", "1GB", "aura", "leipzig").id
     db2 = aura_api.create_instance("db2", "1GB", "aura", "dresden").id
-    aura_api.create_session("one", db1, "12345")
-    aura_api.create_session("other", db2, "123123")
+    aura_api.create_session("one", db1, "12345", memory="8GB")
+    aura_api.create_session("other", db2, "123123", memory="8GB")
 
     sessions = DedicatedSessions(aura_api)
 
@@ -237,7 +237,7 @@ def test_delete_session(aura_api: AuraApi) -> None:
 
 def test_delete_nonexisting_session(aura_api: AuraApi) -> None:
     db1 = aura_api.create_instance("db1", "1gb", "aura", "leipzig").id
-    aura_api.create_session("one", db1, "12345")
+    aura_api.create_session("one", db1, "12345", memory="8GB")
     sessions = DedicatedSessions(aura_api)
 
     assert sessions.delete("other") is False
@@ -246,8 +246,8 @@ def test_delete_nonexisting_session(aura_api: AuraApi) -> None:
 
 def test_delete_nonunique_session(aura_api: AuraApi) -> None:
     db1 = aura_api.create_instance("db1", "1GB", "aura", "leipzig").id
-    aura_api.create_session("one", db1, "12345")
-    aura_api.create_session("one", db1, "12345")
+    aura_api.create_session("one", db1, "12345", memory="8GB")
+    aura_api.create_session("one", db1, "12345", memory="8GB")
     sessions = DedicatedSessions(aura_api)
 
     with pytest.raises(