From de9775a1e01696731fbb7383372de0c2dcaa6fe0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florentin=20D=C3=B6rre?= Date: Wed, 5 Jun 2024 14:28:30 +0200 Subject: [PATCH 1/4] Parse ttl and make host non optional host is already known in the gds-api now on the first call -> always present. ttl is a recently added field --- graphdatascience/session/aura_api_responses.py | 16 ++++++++++++---- graphdatascience/tests/unit/test_aura_api.py | 13 ++++++++++--- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/graphdatascience/session/aura_api_responses.py b/graphdatascience/session/aura_api_responses.py index 008a26cb1..2a38c4c08 100644 --- a/graphdatascience/session/aura_api_responses.py +++ b/graphdatascience/session/aura_api_responses.py @@ -4,8 +4,10 @@ import sys from collections import defaultdict from dataclasses import dataclass -from datetime import datetime, timezone from typing import Any, Dict, NamedTuple, Optional, Set +from datetime import datetime, timedelta, timezone + +from pandas import Timedelta @dataclass(repr=True, frozen=True) @@ -15,13 +17,15 @@ class SessionDetails: instance_id: str memory: str status: str - host: Optional[str] - expiry_date: Optional[datetime] + host: str created_at: datetime + expiry_date: Optional[datetime] + ttl: Optional[timedelta] @classmethod def fromJson(cls, json: Dict[str, Any]) -> SessionDetails: expiry_date = json.get("expiry_date") + ttl = json.get("ttl") return cls( id=json["id"], @@ -29,14 +33,18 @@ def fromJson(cls, json: Dict[str, Any]) -> SessionDetails: instance_id=json["instance_id"], memory=json["memory"], status=json["status"], - host=json.get("host"), + host=json["host"], expiry_date=TimeParser.fromisoformat(expiry_date) if expiry_date else None, created_at=TimeParser.fromisoformat(json["created_at"]), + ttl=Timedelta(ttl).to_pytimedelta() if ttl else None, # datetime has no support for parsing timedetla ) def bolt_connection_url(self) -> str: return f"neo4j+ssc://{self.host}" # TODO use neo4j+s + def is_expired(self) -> bool: + return self.status == "Expired" + @dataclass(repr=True, frozen=True) class InstanceDetails: diff --git a/graphdatascience/tests/unit/test_aura_api.py b/graphdatascience/tests/unit/test_aura_api.py index 061baa8dc..32c55b600 100644 --- a/graphdatascience/tests/unit/test_aura_api.py +++ b/graphdatascience/tests/unit/test_aura_api.py @@ -1,5 +1,5 @@ import logging -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone import pytest from _pytest.logging import LogCaptureFixture @@ -48,6 +48,7 @@ def test_create_session(requests_mock: Mocker) -> None: host="1.2.3.4", memory="4G", expiry_date=None, + ttl=None, ) @@ -82,6 +83,7 @@ def test_list_session(requests_mock: Mocker) -> None: host="1.2.3.4", memory="4G", expiry_date=TimeParser.fromisoformat("1977-01-01T00:00:00Z"), + ttl=None, ) @@ -124,6 +126,7 @@ def test_list_sessions(requests_mock: Mocker) -> None: host="1.2.3.4", memory="4G", expiry_date=TimeParser.fromisoformat("1977-01-01T00:00:00Z"), + ttl=None, ) expected2 = SessionDetails( @@ -133,8 +136,9 @@ def test_list_sessions(requests_mock: Mocker) -> None: instance_id="dbid-3", created_at=TimeParser.fromisoformat("2012-01-01T00:00:00Z"), memory="8G", - host=None, + host="foo.bar", expiry_date=None, + ttl=None, ) assert result == [expected1, expected2] @@ -635,6 +639,7 @@ def test_parse_session_info() -> None: "expiry_date": "2022-01-01T00:00:00Z", "created_at": "2021-01-01T00:00:00Z", "host": "a.b", + "ttl": "1d8h1m2s", } session_info = SessionDetails.fromJson(session_details) @@ -647,10 +652,11 @@ def test_parse_session_info() -> None: host="a.b", expiry_date=datetime(2022, 1, 1, 0, 0, 0, tzinfo=timezone.utc), created_at=datetime(2021, 1, 1, 0, 0, 0, tzinfo=timezone.utc), + ttl=timedelta(days=1, hours=8, minutes=1, seconds=2), ) -def test_parse_session_info_without_expiry() -> None: +def test_parse_session_info_without_optionals() -> None: session_details = { "id": "test_id", "name": "test_session", @@ -670,5 +676,6 @@ def test_parse_session_info_without_expiry() -> None: host="a.b", status="running", expiry_date=None, + ttl=None, created_at=datetime(2021, 1, 1, 0, 0, 0, tzinfo=timezone.utc), ) From 1834540b23181bed2ac459f6698f3d7da215f97b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florentin=20D=C3=B6rre?= Date: Wed, 5 Jun 2024 14:55:35 +0200 Subject: [PATCH 2/4] Check expiry date before connect to session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Max Kießling --- .../session/dedicated_sessions.py | 10 ++++ .../tests/unit/test_dedicated_sessions.py | 59 ++++++++++++++++++- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/graphdatascience/session/dedicated_sessions.py b/graphdatascience/session/dedicated_sessions.py index 37a04b3bb..94260ac08 100644 --- a/graphdatascience/session/dedicated_sessions.py +++ b/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] diff --git a/graphdatascience/tests/unit/test_dedicated_sessions.py b/graphdatascience/tests/unit/test_dedicated_sessions.py index 5747e784e..69653fdd4 100644 --- a/graphdatascience/tests/unit/test_dedicated_sessions.py +++ b/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,6 +53,7 @@ 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 @@ -60,6 +61,12 @@ def create_session(self, name: str, dbid: str, pwd: str, memory: str) -> Session 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 From 24e367eea62e761cbb47f9c752718e073265d50d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florentin=20D=C3=B6rre?= Date: Wed, 5 Jun 2024 16:57:34 +0200 Subject: [PATCH 3/4] Format code --- graphdatascience/session/aura_api_responses.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphdatascience/session/aura_api_responses.py b/graphdatascience/session/aura_api_responses.py index 2a38c4c08..cd533a192 100644 --- a/graphdatascience/session/aura_api_responses.py +++ b/graphdatascience/session/aura_api_responses.py @@ -4,8 +4,8 @@ import sys from collections import defaultdict from dataclasses import dataclass -from typing import Any, Dict, NamedTuple, Optional, Set from datetime import datetime, timedelta, timezone +from typing import Any, Dict, NamedTuple, Optional, Set from pandas import Timedelta From 9afcc263a67cad30f85046f7628c5ec7e8b70c62 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florentin=20D=C3=B6rre?= Date: Wed, 5 Jun 2024 17:44:25 +0200 Subject: [PATCH 4/4] Fix unit test with missing host --- graphdatascience/tests/unit/test_aura_api.py | 1 + 1 file changed, 1 insertion(+) diff --git a/graphdatascience/tests/unit/test_aura_api.py b/graphdatascience/tests/unit/test_aura_api.py index 32c55b600..5b87816ca 100644 --- a/graphdatascience/tests/unit/test_aura_api.py +++ b/graphdatascience/tests/unit/test_aura_api.py @@ -111,6 +111,7 @@ def test_list_sessions(requests_mock: Mocker) -> None: "instance_id": "dbid-3", "created_at": "2012-01-01T00:00:00Z", "memory": "8G", + "host": "foo.bar", }, ], )