Check expiry date before connect to session

Co-authored-by: Max Kießling <max.kiessling@neo4j.com>

Author: FlorentinD

graphdatascience/session/dedicated_sessions.py

@@ -2,6 +2,7 @@
 
 import hashlib
 import warnings
+from datetime import datetime, timedelta, timezone
 from typing import List, Optional
 
 from graphdatascience.session.algorithm_category import AlgorithmCategory
@@ -50,6 +51,7 @@ def get_or_create(
 
         # TODO configure session size (and check existing_session has same size)
         if existing_session:
+            self._check_expiry_date(existing_session)
             session_id = existing_session.id
         else:
             create_details = self._create_session(session_name, dbid, db_connection.uri, password, memory)
@@ -129,6 +131,14 @@ def _construct_client(
             delete_fn=lambda: self.delete(session_name, dbid=AuraApi.extract_id(db_connection.uri)),
         )
 
+    def _check_expiry_date(self, session: SessionDetails) -> None:
+        if session.is_expired():
+            raise RuntimeError(f"Session `{session.name}` is expired. Please delete it and create a new one.")
+        if session.expiry_date:
+            until_expiry: timedelta = session.expiry_date - datetime.now(timezone.utc)
+            if until_expiry < timedelta(days=1):
+                raise Warning(f"Session `{session.name}` is expiring in less than a day.")
+
     @classmethod
     def _fail_ambiguous_session(cls, session_name: str, sessions: List[SessionDetails]) -> None:
         candidates = [i.id for i in sessions]

graphdatascience/tests/unit/test_dedicated_sessions.py

@@ -1,7 +1,7 @@
 import dataclasses
 import re
-from datetime import datetime
-from typing import List, Optional
+from datetime import datetime, timedelta, timezone
+from typing import List, Optional, cast
 
 import pytest
 from pytest_mock import MockerFixture
@@ -53,13 +53,20 @@ def create_session(self, name: str, dbid: str, pwd: str, memory: str) -> Session
             created_at=datetime.fromisoformat("2021-01-01T00:00:00+00:00"),
             host="foo.bar",
             expiry_date=None,
+            ttl=None,
         )
 
         self.id_counter += 1
         self._sessions[details.id] = details
 
         return details
 
+    def add_session(self, session: SessionDetails) -> None:
+        if session.id in self._sessions:
+            raise ValueError(f"Session with id {session.id} already exists.")
+
+        self._sessions[session.id] = session
+
     def create_instance(self, name: str, memory: str, cloud_provider: str, region: str) -> InstanceCreateDetails:
         id = f"ffff{self.id_counter}"
         create_details = InstanceCreateDetails(
@@ -236,6 +243,54 @@ def test_get_or_create_duplicate_session(aura_api: AuraApi) -> None:
         sessions.get_or_create("one", SessionMemory.m_8GB, DbmsConnectionInfo(db.connection_url, "", ""))
 
 
+def test_get_or_create_expired_session(aura_api: AuraApi) -> None:
+    db = _setup_db_instance(aura_api)
+
+    fake_aura_api = cast(FakeAuraApi, aura_api)
+    fake_aura_api.add_session(
+        SessionDetails(
+            id="ffff0-ffff1",
+            name="one",
+            instance_id=db.id,
+            memory=SessionMemory.m_8GB.value,
+            status="Expired",
+            created_at=datetime.now(),
+            host="foo.bar",
+            expiry_date=None,
+            ttl=None,
+        )
+    )
+
+    with pytest.raises(
+        RuntimeError, match=re.escape("Session `one` is expired. Please delete it and create a new one.")
+    ):
+        sessions = DedicatedSessions(aura_api)
+        sessions.get_or_create("one", SessionMemory.m_8GB, DbmsConnectionInfo(db.connection_url, "", ""))
+
+
+def test_get_or_create_soon_expired_session(aura_api: AuraApi) -> None:
+    db = _setup_db_instance(aura_api)
+
+    fake_aura_api = cast(FakeAuraApi, aura_api)
+    fake_aura_api.add_session(
+        SessionDetails(
+            id="ffff0-ffff1",
+            name="one",
+            instance_id=db.id,
+            memory=SessionMemory.m_8GB.value,
+            status="Ready",
+            created_at=datetime.now(),
+            host="foo.bar",
+            expiry_date=datetime.now(tz=timezone.utc) - timedelta(hours=23),
+            ttl=None,
+        )
+    )
+
+    with pytest.raises(Warning, match=re.escape("Session `one` is expiring in less than a day.")):
+        sessions = DedicatedSessions(aura_api)
+        sessions.get_or_create("one", SessionMemory.m_8GB, DbmsConnectionInfo(db.connection_url, "", ""))
+
+
 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