use aura graph analytics

python-wrapper/pyproject.toml

@@ -33,7 +33,8 @@ dependencies = [
     "ipython >=7, <10",
     "pydantic >=2 , <3",
     "pydantic-extra-types >=2, <3",
-    "enum-tools==0.13.0"
+    "enum-tools==0.13.0",
+    "python-dotenv>=1.2.1",
 ]
 requires-python = ">=3.10"
 
@@ -72,7 +73,7 @@ notebook = [
     "palettable>=3.3.3",
     "matplotlib>=3.9.4",
     "snowflake-snowpark-python==1.42.0",
-    "dotenv"
+    "python-dotenv"
 ]
 
 [project.urls]

python-wrapper/tests/conftest.py

@@ -1,4 +1,5 @@
 import os
+import random
 from typing import Any, Generator
 
 import pytest
@@ -31,45 +32,56 @@ def pytest_collection_modifyitems(config: Any, items: Any) -> None:
 
 
 @pytest.fixture(scope="package")
-def aura_ds_instance() -> Generator[Any, None, None]:
+def aura_db_instance() -> Generator[Any, None, None]:
     if os.environ.get("AURA_API_CLIENT_ID", None) is None:
         yield None
         return
 
-    from tests.gds_helper import aura_api, create_aurads_instance
+    from tests.gds_helper import aura_api, create_auradb_instance
 
     api = aura_api()
-    id, dbms_connection_info = create_aurads_instance(api)
+    dbms_connection_info = create_auradb_instance(api)
 
+    old_uri = os.environ.get("NEO4J_URI", "")
     # setting as environment variables to run notebooks with this connection
     os.environ["NEO4J_URI"] = dbms_connection_info.get_uri()
     assert isinstance(dbms_connection_info.username, str)
     os.environ["NEO4J_USER"] = dbms_connection_info.username
     assert isinstance(dbms_connection_info.password, str)
     os.environ["NEO4J_PASSWORD"] = dbms_connection_info.password
+    old_instance = os.environ.get("AURA_INSTANCEID", "")
+    if dbms_connection_info.aura_instance_id:
+        os.environ["AURA_INSTANCEID"] = dbms_connection_info.aura_instance_id
+
     yield dbms_connection_info
 
     # Clear Neo4j_URI after test (rerun should create a new instance)
-    os.environ["NEO4J_URI"] = ""
-    api.delete_instance(id)
+    os.environ["NEO4J_URI"] = old_uri
+    os.environ["AURA_INSTANCEID"] = old_instance
+    assert dbms_connection_info.aura_instance_id is not None
+    api.delete_instance(dbms_connection_info.aura_instance_id)
 
 
 @pytest.fixture(scope="package")
-def gds(aura_ds_instance: Any) -> Generator[Any, None, None]:
-    from graphdatascience import GraphDataScience
+def gds(aura_db_instance: Any) -> Generator[Any, None, None]:
+    from graphdatascience.session import SessionMemory
 
-    from tests.gds_helper import connect_to_plugin_gds
+    from tests.gds_helper import connect_to_plugin_gds, gds_sessions
 
-    if aura_ds_instance:
-        yield GraphDataScience(
-            endpoint=aura_ds_instance.uri,
-            auth=(aura_ds_instance.username, aura_ds_instance.password),
-            aura_ds=True,
-            database="neo4j",
+    if aura_db_instance:
+        sessions = gds_sessions()
+
+        gds = sessions.get_or_create(
+            f"neo4j-viz-ci-{os.environ.get('GITHUB_RUN_ID', random.randint(0, 10**6))}",
+            memory=SessionMemory.m_2GB,
+            db_connection=aura_db_instance,
         )
+
+        yield gds
+        gds.delete()
     else:
-        NEO4J_URI = os.environ.get("NEO4J_URI", "neo4j://localhost:7687")
-        gds = connect_to_plugin_gds(NEO4J_URI)
+        NEO4J_URI = os.environ["NEO4J_URI"]
+        gds = connect_to_plugin_gds(NEO4J_URI)  # type: ignore
         yield gds
         gds.close()
 

python-wrapper/tests/gds_helper.py

@@ -1,9 +1,9 @@
 import os
 import re
 
-from graphdatascience import GraphDataScience
+from graphdatascience import GdsSessions, GraphDataScience
 from graphdatascience.semantic_version.semantic_version import SemanticVersion
-from graphdatascience.session import DbmsConnectionInfo, SessionMemory
+from graphdatascience.session import AuraAPICredentials, DbmsConnectionInfo, SessionMemory
 from graphdatascience.session.aura_api import AuraApi
 from graphdatascience.session.aura_api_responses import InstanceCreateDetails
 from graphdatascience.version import __version__
@@ -49,21 +49,29 @@ def aura_api() -> AuraApi:
         )
 
 
-def create_aurads_instance(api: AuraApi) -> tuple[str, DbmsConnectionInfo]:
-    # Switch to Sessions once they can be created without a DB
+def gds_sessions() -> GdsSessions:
+    return GdsSessions(
+        api_credentials=AuraAPICredentials(
+            client_id=os.environ["AURA_API_CLIENT_ID"],
+            client_secret=os.environ["AURA_API_CLIENT_SECRET"],
+            project_id=os.environ.get("AURA_API_TENANT_ID"),
+        )
+    )
+
+
+def create_auradb_instance(api: AuraApi) -> DbmsConnectionInfo:
     instance_details: InstanceCreateDetails = api.create_instance(
-        name="ci-neo4j-viz-session",
-        memory=SessionMemory.m_8GB.value,
+        name="ci-neo4j-viz-db",
+        memory=SessionMemory.m_2GB.value,
         cloud_provider="gcp",
         region="europe-west1",
+        type="enterprise-db",
     )
 
     wait_result = api.wait_for_instance_running(instance_id=instance_details.id)
     if wait_result.error:
         raise Exception(f"Error while waiting for instance to be running: {wait_result.error}")
 
-    return instance_details.id, DbmsConnectionInfo(
-        uri=wait_result.connection_url,
-        username="neo4j",
-        password=instance_details.password,
+    return DbmsConnectionInfo(
+        username="neo4j", password=instance_details.password, aura_instance_id=instance_details.id
     )