Handle reconnect to non-running instances

Before we got into odd state for instances with state creating or destroying

Author: FlorentinD

graphdatascience/gds_session/aura_api.py

@@ -6,7 +6,7 @@
 import time
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, List, Optional, Set
+from typing import Any, List, NamedTuple, Optional, Set
 from urllib.parse import urlparse
 
 import requests as req
@@ -71,6 +71,19 @@ def from_json(cls, json: dict[str, Any]) -> 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
@@ -210,16 +223,16 @@ def list_instance(self, instance_id: str) -> Optional[InstanceSpecificDetails]:
 
     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. "
@@ -229,7 +242,7 @@ def wait_for_instance_running(
             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(

graphdatascience/gds_session/gds_sessions.py

@@ -67,11 +67,11 @@ def get_or_create(
             GdsSessions._instance_name(session_name), size.value, db_instance.cloud_provider, region
         )
         wait_result = self._aura_api.wait_for_instance_running(create_details.id)
-        if wait_result is not None:
-            raise RuntimeError(f"Failed to create session `{session_name}`: {wait_result}")
+        if err := wait_result.error:
+            raise RuntimeError(f"Failed to create session `{session_name}`: {err}")
 
         gds_user = create_details.username
-        gds_url = create_details.connection_url
+        gds_url = wait_result.connection_url
 
         self._change_initial_pw(
             gds_url=gds_url, gds_user=gds_user, initial_pw=create_details.password, new_pw=db_connection.password
@@ -125,14 +125,10 @@ def _try_connect(self, session_name: str, db_connection: DbmsConnectionInfo) ->
         if len(matched_instances) > 1:
             self._fail_ambiguous_session(session_name, matched_instances)
 
-        instance_details = self._aura_api.list_instance(matched_instances[0].id)
-
-        if instance_details:
-            gds_url = instance_details.connection_url
-        else:
-            raise RuntimeError(
-                f"Unable to get connection information for session `{session_name}`. Does it still exist?"
-            )
+        wait_result = self._aura_api.wait_for_instance_running(matched_instances[0].id)
+        if err := wait_result.error:
+            raise RuntimeError(f"Failed to connect to session `{session_name}`: {err}")
+        gds_url = wait_result.connection_url
 
         return self._construct_client(session_name=session_name, gds_url=gds_url, db_connection=db_connection)
 

graphdatascience/tests/unit/test_aura_api.py

@@ -10,6 +10,7 @@
     InstanceCreateDetails,
     InstanceSpecificDetails,
     TenantDetails,
+    WaitResult,
 )
 
 
@@ -88,10 +89,12 @@ def test_create_instance(requests_mock: Mocker) -> None:
 
     requests_mock.get(
         "https://api.neo4j.io/v1/tenants/some-tenant",
-        json={"data": {
-            "id": "some_tenant",
-            "instance_configurations": [{"type": "enterprise-ds", "region": "leipzig-1", "cloud_provider": "aws"}]
-            }},
+        json={
+            "data": {
+                "id": "some_tenant",
+                "instance_configurations": [{"type": "enterprise-ds", "region": "leipzig-1", "cloud_provider": "aws"}],
+            }
+        },
     )
 
     requests_mock.post(
@@ -264,8 +267,9 @@ def test_dont_wait_forever(requests_mock: Mocker, caplog: LogCaptureFixture) ->
     api = AuraApi("", "", tenant_id="some-tenant")
 
     with caplog.at_level(logging.DEBUG):
-        assert "Instance is not running after waiting for 0.8" in api.wait_for_instance_running(  # type: ignore
-            "id0", max_sleep_time=0.7
+        assert (
+            "Instance is not running after waiting for 0.8"
+            in api.wait_for_instance_running("id0", max_sleep_time=0.7).error
         )
 
     assert "Instance `id0` is not yet running. Current status: creating. Retrying in 0.2 seconds..." in caplog.text
@@ -279,7 +283,7 @@ def test_wait_for_instance_running(requests_mock: Mocker) -> None:
             "data": {
                 "status": "running",
                 "cloud_provider": None,
-                "connection_url": None,
+                "connection_url": "foo.bar",
                 "id": None,
                 "name": None,
                 "region": None,
@@ -291,7 +295,7 @@ def test_wait_for_instance_running(requests_mock: Mocker) -> None:
 
     api = AuraApi("", "", tenant_id="some-tenant")
 
-    assert api.wait_for_instance_running("id0") is None
+    assert api.wait_for_instance_running("id0") == WaitResult.from_connection_url("foo.bar")
 
 
 def test_wait_for_instance_deleting(requests_mock: Mocker) -> None:
@@ -329,8 +333,8 @@ def test_wait_for_instance_deleting(requests_mock: Mocker) -> None:
 
     api = AuraApi("", "", tenant_id="some-tenant")
 
-    assert "Instance is being deleted" in api.wait_for_instance_running("id0")  # type: ignore
-    assert "Instance is being deleted" in api.wait_for_instance_running("id1")  # type: ignore
+    assert api.wait_for_instance_running("id0") == WaitResult.from_error("Instance is being deleted")
+    assert api.wait_for_instance_running("id1") == WaitResult.from_error("Instance is being deleted")
 
 
 def test_extract_id() -> None:

graphdatascience/tests/unit/test_gds_sessions.py

@@ -13,6 +13,7 @@
     InstanceDetails,
     InstanceSpecificDetails,
     TenantDetails,
+    WaitResult,
 )
 from graphdatascience.gds_session.dbms_connection_info import DbmsConnectionInfo
 from graphdatascience.gds_session.gds_sessions import (
@@ -78,7 +79,7 @@ def list_instance(self, instance_id: str) -> Optional[InstanceSpecificDetails]:
 
     def wait_for_instance_running(
         self, instance_id: str, sleep_time: float = 0.2, max_sleep_time: float = 300
-    ) -> Optional[str]:
+    ) -> WaitResult:
         return super().wait_for_instance_running(instance_id, sleep_time=0.0001, max_sleep_time=0.001)
 
     def tenant_details(self) -> TenantDetails:
@@ -406,7 +407,8 @@ def test_create_session_invalid_region(aura_api: AuraApi) -> None:
     sessions._aura_api = aura_api
 
     expected_message = (
-        "Region `only-db-region` is not supported by the tenant `tenant_id`." " Supported regions: {'leipzig-1', 'dresden-2'}."
+        "Region `only-db-region` is not supported by the tenant `tenant_id`."
+        " Supported regions: {'leipzig-1', 'dresden-2'}."
     )
     with pytest.raises(ValueError, match=expected_message):
         sessions.get_or_create(