Merge pull request #674 from FlorentinD/ignore-paused-dbs

Skip paused dbs for sessions.list()

Author: FlorentinD

graphdatascience/query_runner/gds_arrow_client.py

@@ -32,10 +32,14 @@ def create(
         tls_root_certs: Optional[bytes] = None,
         connection_string_override: Optional[str] = None,
     ) -> "GdsArrowClient":
+        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=["listenAddress", "versions"]
-            )
+            query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False, yields=debugYields)
             .squeeze()
             .to_dict()
         )

graphdatascience/session/dedicated_sessions.py

@@ -5,6 +5,8 @@
 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
@@ -94,7 +96,12 @@ def list(self) -> List[SessionInfo]:
 
         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]
 

graphdatascience/tests/unit/test_dedicated_sessions.py

@@ -5,6 +5,8 @@
 
 import pytest
 from pytest_mock import MockerFixture
+from requests import Response
+from requests.exceptions import HTTPError
 
 from graphdatascience.session.algorithm_category import AlgorithmCategory
 from graphdatascience.session.aura_api import AuraApi
@@ -102,6 +104,14 @@ def delete_instance(self, instance_id: str) -> InstanceSpecificDetails:
         return self._instances.pop(instance_id)
 
     def list_sessions(self, dbid: str) -> List[SessionDetails]:
+        # mimic aura behaviour of paused instances not being in an orchestra
+        db = self.list_instance(dbid)
+        if db and db.status == "paused":
+            response = Response()
+            response.status_code = 404
+            response._content = b"database not found"
+            raise HTTPError(request=None, response=response)
+
         return [v for _, v in self._sessions.items() if v.instance_id == dbid]
 
     def list_instances(self) -> List[InstanceDetails]:
@@ -180,6 +190,35 @@ def test_list_session(aura_api: AuraApi) -> None:
     assert sessions.list() == [SessionInfo.from_session_details(session)]
 
 
+def test_list_session_paused_instance(aura_api: AuraApi) -> None:
+    db = aura_api.create_instance("test", SessionMemory.m_8GB.value, "aws", "leipzig-1")
+    fake_aura_api = cast(FakeAuraApi, aura_api)
+
+    fake_aura_api.id_counter += 1
+    paused_db = InstanceSpecificDetails(
+        id="4242",
+        status="paused",
+        connection_url="foo.bar",
+        memory=SessionMemory.m_16GB.value,
+        type="",
+        region="dresden",
+        name="paused-db",
+        tenant_id=fake_aura_api._tenant_id,
+        cloud_provider="aws",
+    )
+    fake_aura_api._instances[paused_db.id] = paused_db
+
+    session = aura_api.create_session(
+        name="gds-session-my-session-name",
+        dbid=db.id,
+        pwd="some_pwd",
+        memory=SessionMemory.m_8GB.value,
+    )
+    sessions = DedicatedSessions(aura_api)
+
+    assert sessions.list() == [SessionInfo.from_session_details(session)]
+
+
 def test_create_session(mocker: MockerFixture, aura_api: AuraApi) -> None:
     _setup_db_instance(aura_api)