Check db env through cypher and allow aura_instance_id to be filled from env

Author: FlorentinD

doc/modules/ROOT/pages/graph-analytics-serverless.adoc

@@ -98,10 +98,9 @@ sessions.get_or_create(
 | Name                        | Type               | Optional | Default | Description
 | session_name                | str                | no       | -       | Name of the session. Must be unique within the project.
 | memory                      | https://neo4j.com/docs/graph-data-science-client/{docs-version}/api/sessions/session_memory[SessionMemory]      | no       | -       | Amount of memory available to the session.
-| db_connection               |  https://neo4j.com/docs/graph-data-science-client/{docs-version}/api/sessions/dbms_connection_info[DbmsConnectionInfo] | yes      | None    | Bolt server URL, username, and password to a Neo4j DBMS. Required for the Attached and Self-managed types. Alternatively to username and password, you can provide a `neo4j.Auth` https://neo4j.com/docs/python-manual/current/connect-advanced/#authentication-methods[object].
+| db_connection               |  https://neo4j.com/docs/graph-data-science-client/{docs-version}/api/sessions/dbms_connection_info[DbmsConnectionInfo] | yes      | None    | Bolt server URI, username, and password to a Neo4j DBMS. Required for the Attached and Self-managed types. Alternatively to username and password, you can provide a `neo4j.Auth` https://neo4j.com/docs/python-manual/current/connect-advanced/#authentication-methods[object].
 | ttl                         | datetime.timedelta | yes      | 1h      | Time-to-live for the session.
 | cloud_location              | https://neo4j.com/docs/graph-data-science-client/{docs-version}/api/sessions/cloud_location[CloudLocation]      | yes      | None    | Aura-supported cloud provider and region where the GDS Session will run. Required for the Self-managed and Standalone types.
-| aura_instance_id            | str                | yes      | None    | Aura instance ID for the session. Required for the Attached type, if the id could not be derived from the DB connection.
 | timeout                     | int                | yes      | None    | Seconds to wait for the session to enter Ready state. If the time is exceeded, an error will be returned.
 | neo4j_driver_options        | dict[str, any]     | yes      | None    | Additional options passed to the Neo4j driver to the Neo4j DBMS. Only relevant if `db_connection` is specified.
 | arrow_client_options        | dict[str, any]     | yes      | None    | Additional options passed to the Arrow Flight Client used to connect to the Session.
@@ -317,7 +316,13 @@ from graphdatascience.session import SessionMemory, DbmsConnectionInfo, GdsSessi
 
 sessions = GdsSessions(api_credentials=AuraAPICredentials(os.environ["CLIENT_ID"], os.environ["CLIENT_SECRET"]))
 
-db_connection = DbmsConnectionInfo(os.environ["DB_URI"], os.environ["DB_USER"], os.environ["DB_PASSWORD"])
+# you can also use DbmsConnectionInfo.from_env() to load credentials from environment variables
+db_connection = DbmsConnectionInfo(
+    uri=os.environ["NEO4J_URI"],
+    username=os.environ["NEO4J_USERNAME"],
+    password=os.environ["NEO4J_PASSWORD"],
+    aura_instance_id=os.environ["AURA_INSTANCEID"]
+)
 gds = sessions.get_or_create(
     session_name="my-new-session",
     memory=SessionMemory.m_8GB,

graphdatascience/query_runner/db_environment_resolver.py

@@ -0,0 +1,17 @@
+from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
+
+
+class DbEnvironmentResolver:
+    @staticmethod
+    def hosted_in_aura(db_runner: Neo4jQueryRunner) -> bool:
+        return (
+            db_runner.run_retryable_cypher("""
+        CALL dbms.components() YIELD name, versions
+        WHERE name = "Neo4j Kernel"
+        UNWIND versions as v
+        WITH name, v
+        WHERE v ENDS WITH "aura"
+        RETURN count(*) <> 0
+        """).squeeze()
+            is True
+        )

graphdatascience/session/dbms_connection_info.py

@@ -2,7 +2,6 @@
 
 import os
 from dataclasses import dataclass
-from urllib.parse import urlparse
 
 from neo4j import Auth, basic_auth
 
@@ -20,6 +19,7 @@ class DbmsConnectionInfo:
     database: str | None = None
     # Optional: typed authentication, used instead of username/password. Supports for example a token. See https://neo4j.com/docs/python-manual/current/connect-advanced/#authentication-methods
     auth: Auth | None = None
+    aura_instance_id: str | None = None
 
     def __post_init__(self) -> None:
         # Validate auth fields
@@ -39,14 +39,6 @@ def get_auth(self) -> Auth | None:
             auth = basic_auth(self.username, self.password)
         return auth
 
-    def hosted_in_aura(self) -> bool:
-        """
-        Returns:
-            True if the DBMS is hosted in Aura, False otherwise.
-        """
-        host = urlparse(self.uri).hostname
-        return host is not None and host.endswith("databases.neo4j.io")
-
     @staticmethod
     def from_env() -> DbmsConnectionInfo:
         """
@@ -56,10 +48,12 @@ def from_env() -> DbmsConnectionInfo:
         - NEO4J_USERNAME
         - NEO4J_PASSWORD
         - NEO4J_DATABASE
+        - AURA_INSTANCEID
         """
         uri = os.environ["NEO4J_URI"]
         username = os.environ.get("NEO4J_USERNAME", "neo4j")
         password = os.environ["NEO4J_PASSWORD"]
         database = os.environ.get("NEO4J_DATABASE")
+        aura_instance_id = os.environ.get("AURA_INSTANCEID")
 
-        return DbmsConnectionInfo(uri, username, password, database)
+        return DbmsConnectionInfo(uri, username, password, database, aura_instance_id=aura_instance_id)

graphdatascience/session/dedicated_sessions.py

@@ -6,6 +6,7 @@
 from typing import Any
 
 from graphdatascience.query_runner.arrow_authentication import ArrowAuthentication
+from graphdatascience.query_runner.db_environment_resolver import DbEnvironmentResolver
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 from graphdatascience.session.algorithm_category import AlgorithmCategory
 from graphdatascience.session.aura_api import AuraApi
@@ -61,7 +62,6 @@ def get_or_create(
         db_connection: DbmsConnectionInfo | None = None,
         ttl: timedelta | None = None,
         cloud_location: CloudLocation | None = None,
-        aura_instance_id: str | None = None,
         timeout: int | None = None,
         neo4j_driver_options: dict[str, Any] | None = None,
         arrow_client_options: dict[str, Any] | None = None,
@@ -75,17 +75,21 @@ def get_or_create(
         else:
             db_runner = self._create_db_runner(db_connection, neo4j_driver_options)
 
-            aura_instance_id = AuraApi.extract_id(db_connection.uri) if not aura_instance_id else aura_instance_id
+            aura_instance_id = (
+                db_connection.aura_instance_id
+                if db_connection.aura_instance_id
+                else AuraApi.extract_id(db_connection.uri)
+            )
+            if not aura_instance_id and DbEnvironmentResolver.hosted_in_aura(db_runner):
+                raise ValueError(
+                    f"Could not derive Aura instance id from the URI `{db_connection.uri}`. Please specify the `aura_instance_id` in the `db_connection` argument."
+                )
+
             aura_db_instance = self._aura_api.list_instance(aura_instance_id)
 
             if aura_db_instance is None:
                 if not cloud_location:
-                    if db_connection.hosted_in_aura():
-                        raise ValueError(
-                            f"Could not derive Aura instance id from the URI `{db_connection.uri}`. Please provide the instance id via the `aura_instance_id` argument, or specify a cloud location if the DBMS is self-managed."
-                        )
-                    else:
-                        raise ValueError("cloud_location must be provided for sessions against a self-managed DB.")
+                    raise ValueError("cloud_location must be provided for sessions against a self-managed DB.")
 
                 session_details = self._get_or_create_self_managed_session(
                     session_name, memory.value, cloud_location, ttl

graphdatascience/session/gds_sessions.py

@@ -107,7 +107,6 @@ def get_or_create(
         timeout: int | None = None,
         neo4j_driver_config: dict[str, Any] | None = None,
         arrow_client_options: dict[str, Any] | None = None,
-        aura_instance_id: str | None = None,
     ) -> AuraGraphDataScience:
         """
         Retrieves an existing session with the given session name and database connection,
@@ -125,7 +124,6 @@ def get_or_create(
             timeout (int | None): Optional timeout (in seconds) when waiting for session to become ready. If unset the method will wait forever. If set and session does not become ready an exception will be raised. It is user responsibility to ensure resource gets cleaned up in this situation.
             neo4j_driver_config (dict[str, Any] | None): Optional configuration for the Neo4j driver to the Neo4j DBMS. Only relevant if `db_connection` is specified..
             arrow_client_options (dict[str, Any] | None): Optional configuration for the Arrow Flight client.
-            aura_instance_id (str | None): The Aura instance id. Required if the database is in Aura but its instance id cannot be inferred from the connection information.
         Returns:
             AuraGraphDataScience: The session.
         """
@@ -134,7 +132,6 @@ def get_or_create(
             memory,
             db_connection=db_connection,
             ttl=ttl,
-            aura_instance_id=aura_instance_id,
             cloud_location=cloud_location,
             timeout=timeout,
             neo4j_driver_options=neo4j_driver_config,

graphdatascience/tests/integration/test_db_environment_resolver.py

@@ -0,0 +1,14 @@
+import pytest
+
+from graphdatascience.query_runner.db_environment_resolver import DbEnvironmentResolver
+from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
+
+
+@pytest.mark.only_on_aura
+def test_hosted_in_aura_aura_dbms(aura_runner: Neo4jQueryRunner) -> None:
+    assert DbEnvironmentResolver.hosted_in_aura(aura_runner)
+
+
+@pytest.mark.skip_on_aura
+def test_hosted_in_aura_self_managed_dbms(runner: Neo4jQueryRunner) -> None:
+    assert not DbEnvironmentResolver.hosted_in_aura(runner)

graphdatascience/tests/unit/session/test_dbms_connection_info.py

@@ -39,10 +39,3 @@ def test_dbms_connection_info_fail_on_auth_and_username() -> None:
         )
     else:
         assert False, "Expected ValueError was not raised"
-
-
-def test_dbms_connection_info_hosted_in_aura() -> None:
-    assert DbmsConnectionInfo(uri="bolt://something.databases.neo4j.io").hosted_in_aura()
-    assert DbmsConnectionInfo(uri="bolt://something.databases.neo4j.io:7474").hosted_in_aura()
-
-    assert not DbmsConnectionInfo(uri="bolt://something.neo4j.com").hosted_in_aura()

graphdatascience/tests/unit/test_dedicated_sessions.py

@@ -431,7 +431,7 @@ def test_create_standalone_session(mocker: MockerFixture, aura_api: AuraApi) ->
     sessions = DedicatedSessions(aura_api)
 
     patch_construct_client(mocker)
-    patch_neo4j_query_runner(mocker)
+    patch_neo4j_query_runner(mocker, hosted_in_aura=False)
 
     ttl = timedelta(hours=42)
 
@@ -526,6 +526,22 @@ def test_get_or_create(mocker: MockerFixture, aura_api: AuraApi) -> None:
     assert [i.name for i in sessions.list()] == ["my-session"]
 
 
+def test_get_or_create_with_explicit_aura_instance_id(mocker: MockerFixture, aura_api: AuraApi) -> None:
+    db = _setup_db_instance(aura_api)
+    sessions = DedicatedSessions(aura_api)
+    patch_construct_client(mocker)
+    patch_neo4j_query_runner(mocker)
+
+    sessions.get_or_create(
+        "my-session",
+        SessionMemory.m_8GB,
+        DbmsConnectionInfo(
+            "neo4j+s://foo.bar", "dbuser", "db_pw", aura_instance_id=db.id
+        ),  # not part of list instances result
+        cloud_location=None,
+    )
+
+
 def test_get_or_create_expired_session(mocker: MockerFixture, aura_api: AuraApi) -> None:
     db = _setup_db_instance(aura_api)
 
@@ -601,7 +617,7 @@ def test_get_or_create_for_auradb_with_cloud_location(mocker: MockerFixture, aur
 
 def test_get_or_create_for_without_cloud_location(mocker: MockerFixture, aura_api: AuraApi) -> None:
     sessions = DedicatedSessions(aura_api)
-    patch_neo4j_query_runner(mocker)
+    patch_neo4j_query_runner(mocker, hosted_in_aura=False)
 
     with pytest.raises(
         ValueError, match=re.escape("cloud_location must be provided for sessions against a self-managed DB.")
@@ -621,7 +637,7 @@ def test_get_or_create_for_non_derivable_aura_instance_id(mocker: MockerFixture,
     with pytest.raises(
         ValueError,
         match=re.escape(
-            "Could not derive Aura instance id from the URI `neo4j+s://06cba79f.databases.neo4j.io`. Please provide the instance id via the `aura_instance_id` argument, or specify a cloud location if the DBMS is self-managed."
+            "Could not derive Aura instance id from the URI `neo4j+s://06cba79f.databases.neo4j.io`. Please specify the `aura_instance_id` in the `db_connection` argument."
         ),
     ):
         sessions.get_or_create(
@@ -808,12 +824,16 @@ def _setup_db_instance(aura_api: AuraApi) -> InstanceCreateDetails:
     return aura_api.create_instance("test", SessionMemory.m_8GB.value, "aws", "leipzig-1")
 
 
-def patch_neo4j_query_runner(mocker: MockerFixture) -> None:
+def patch_neo4j_query_runner(mocker: MockerFixture, hosted_in_aura: bool = True) -> None:
     mocker.patch(
         "graphdatascience.query_runner.neo4j_query_runner.Neo4jQueryRunner.create_for_db",
         lambda *args, **kwargs: kwargs,
     )
     mocker.patch("graphdatascience.session.dedicated_sessions.DedicatedSessions._validate_db_connection")
+    mocker.patch(
+        "graphdatascience.query_runner.db_environment_resolver.DbEnvironmentResolver.hosted_in_aura",
+        lambda *args, **kwargs: hosted_in_aura,
+    )
 
 
 def patch_construct_client(mocker: MockerFixture) -> None: