Validate region and allow custom override

Author: FlorentinD

graphdatascience/gds_session/gds_sessions.py

@@ -45,7 +45,11 @@ def __init__(self, ds_connection: AuraAPICredentials) -> None:
         )
 
     def get_or_create(
-        self, session_name: str, size: SessionSizeByMemory, db_connection: DbmsConnectionInfo
+        self,
+        session_name: str,
+        size: SessionSizeByMemory,
+        db_connection: DbmsConnectionInfo,
+        region: Optional[str] = None,
     ) -> AuraGraphDataScience:
         connected_instance = self._try_connect(session_name, db_connection)
         if connected_instance is not None:
@@ -56,8 +60,11 @@ def get_or_create(
         if not db_instance:
             raise ValueError(f"Could not find Aura instance with the uri `{db_connection.uri}`")
 
+        region = region if region else db_instance.region
+        self._validate_region(region, db_instance)
+
         create_details = self._aura_api.create_instance(
-            GdsSessions._instance_name(session_name), size.value, db_instance.cloud_provider, db_instance.region
+            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:
@@ -129,6 +136,15 @@ def _try_connect(self, session_name: str, db_connection: DbmsConnectionInfo) ->
 
         return self._construct_client(session_name=session_name, gds_url=gds_url, db_connection=db_connection)
 
+    def _validate_region(self, region: str, db_instance: InstanceSpecificDetails) -> None:
+        tenant_details = self._aura_api.tenant_details()
+        available_regions = tenant_details.regions_per_provider[db_instance.cloud_provider]
+        if region not in available_regions:
+            raise ValueError(
+                f"Region `{region}` is not supported by the tenant `{tenant_details.id}`."
+                f" Supported regions: {available_regions}`"
+            )
+
     def _construct_client(
         self, session_name: str, gds_url: str, db_connection: DbmsConnectionInfo
     ) -> AuraGraphDataScience:

graphdatascience/tests/unit/test_gds_sessions.py

@@ -12,6 +12,7 @@
     InstanceCreateDetails,
     InstanceDetails,
     InstanceSpecificDetails,
+    TenantDetails,
 )
 from graphdatascience.gds_session.dbms_connection_info import DbmsConnectionInfo
 from graphdatascience.gds_session.gds_sessions import (
@@ -80,6 +81,11 @@ def wait_for_instance_running(
     ) -> Optional[str]:
         return super().wait_for_instance_running(instance_id, sleep_time=0.0001, max_sleep_time=0.001)
 
+    def tenant_details(self) -> TenantDetails:
+        return TenantDetails(
+            id=self._tenant_id, ds_type="fake-ds", regions_per_provider={"aws": {"leipzig-1", "dresden-2"}}
+        )
+
 
 @pytest.fixture
 def aura_api() -> AuraApi:
@@ -176,6 +182,31 @@ def test_create_default_session(mocker: MockerFixture, aura_api: AuraApi) -> Non
     assert instance_details.memory == "8GB"
 
 
+def test_create_session_override_region(mocker: MockerFixture, aura_api: AuraApi) -> None:
+    _setup_db_instance(aura_api)
+
+    sessions = GdsSessions(AuraAPICredentials("", "", "placeholder"))
+    sessions._aura_api = aura_api
+
+    mocker.patch(
+        "graphdatascience.gds_session.gds_sessions.GdsSessions._construct_client", lambda *args, **kwargs: kwargs
+    )
+    mocker.patch(
+        "graphdatascience.gds_session.gds_sessions.GdsSessions._change_initial_pw", lambda *args, **kwargs: kwargs
+    )
+
+    sessions.get_or_create(
+        "my-session",
+        SessionSizeByMemory.DEFAULT,
+        DbmsConnectionInfo("neo4j+ssc://ffff0.databases.neo4j.io", "dbuser", "db_pw"),
+        region="dresden-2",
+    )
+    instance_details: InstanceSpecificDetails = aura_api.list_instance("ffff1")  # type: ignore
+    assert instance_details.cloud_provider == "aws"
+    assert instance_details.region == "dresden-2"
+    assert instance_details.memory == "8GB"
+
+
 def test_get_or_create(mocker: MockerFixture, aura_api: AuraApi) -> None:
     _setup_db_instance(aura_api)
 
@@ -368,5 +399,22 @@ def test_create_waiting_forever() -> None:
         )
 
 
+def test_create_session_invalid_region(aura_api: AuraApi) -> None:
+    aura_api.create_instance("test", "8GB", "aws", "only-db-region")
+
+    sessions = GdsSessions(AuraAPICredentials("", "", "placeholder"))
+    sessions._aura_api = aura_api
+
+    expected_message = (
+        "Region `only-db-region` is not supported by the tenant `tenant_id`." " Supported regions: {'leipzig-1'}."
+    )
+    with pytest.raises(ValueError, match=expected_message):
+        sessions.get_or_create(
+            "my-session",
+            SessionSizes.by_memory().X5L,
+            DbmsConnectionInfo("neo4j+ssc://ffff0.databases.neo4j.io", "dbuser", "db_pw"),
+        )
+
+
 def _setup_db_instance(aura_api: AuraApi) -> None:
     aura_api.create_instance("test", "8GB", "aws", "leipzig-1")