Merge pull request #549 from FlorentinD/sessions-block-non-remote-projections

Block local projections in AuraGDS

Author: FlorentinD

examples/dev/aura-only-features.ipynb

@@ -316,7 +316,7 @@
     "There are two key differences between the remote projection and Cypher projections V2:\n",
     "\n",
     "1. In AuraDB, the aggregating function does not take a graph name as a parameter.\n",
-    "2. The aggregation function should only be called through the GDS Python Client endpoint `gds.graph.project.remoteDb`\n",
+    "2. The aggregation function should only be called through the GDS Python Client endpoint `gds.graph.project`\n",
     "\n",
     "### Limitations\n",
     "\n",
@@ -335,7 +335,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "G, result = gds.graph.project.remoteDb(\n",
+    "G, result = gds.graph.project(\n",
     "    \"pagerank-graph\",\n",
     "    \"\"\"\n",
     "    MATCH (u:User) \n",

graphdatascience/aura_graph_data_science.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple, Union
 
 from neo4j import GraphDatabase
 from pandas import DataFrame
@@ -8,11 +8,13 @@
 from graphdatascience.call_builder import IndirectCallBuilder
 from graphdatascience.endpoints import AlphaEndpoints, BetaEndpoints, DirectEndpoints
 from graphdatascience.error.uncallable_namespace import UncallableNamespace
+from graphdatascience.graph.graph_proc_runner import GraphRemoteProcRunner
 from graphdatascience.query_runner.arrow_query_runner import ArrowQueryRunner
 from graphdatascience.query_runner.aura_db_arrow_query_runner import (
     AuraDbArrowQueryRunner,
     AuraDbConnectionInfo,
 )
+from graphdatascience.query_runner.query_runner import QueryRunner
 
 
 class AuraGraphDataScience(DirectEndpoints, UncallableNamespace):
@@ -23,46 +25,51 @@ class AuraGraphDataScience(DirectEndpoints, UncallableNamespace):
 
     def __init__(
         self,
-        endpoint: str,
+        endpoint: Union[str, QueryRunner],
         auth: Tuple[str, str],
         aura_db_connection_info: AuraDbConnectionInfo,
         arrow_disable_server_verification: bool = True,
         arrow_tls_root_certs: Optional[bytes] = None,
         bookmarks: Optional[Any] = None,
     ):
-        gds_query_runner = ArrowQueryRunner.create(
-            Neo4jQueryRunner.create(endpoint, auth, aura_ds=True),
-            auth,
-            True,
-            arrow_disable_server_verification,
-            arrow_tls_root_certs,
-        )
-
-        self._server_version = gds_query_runner.server_version()
-
-        if self._server_version < ServerVersion(2, 6, 0):
-            raise RuntimeError(
-                f"AuraDB connection info was provided but GDS version {self._server_version} \
-                    does not support connecting to AuraDB"
+        if isinstance(endpoint, str):
+            gds_query_runner = ArrowQueryRunner.create(
+                Neo4jQueryRunner.create(endpoint, auth, aura_ds=True),
+                auth,
+                True,
+                arrow_disable_server_verification,
+                arrow_tls_root_certs,
             )
 
-        self._driver_config = gds_query_runner.driver_config()
-        driver = GraphDatabase.driver(
-            aura_db_connection_info.uri, auth=aura_db_connection_info.auth, **self._driver_config
-        )
-        self._db_query_runner = Neo4jQueryRunner(
-            driver, auto_close=True, bookmarks=bookmarks, server_version=self._server_version
-        )
+            self._server_version = gds_query_runner.server_version()
 
-        # we need to explicitly set these as the default value is None
-        # which signals the driver to use the default configured database
-        # from the dbms.
-        gds_query_runner.set_database("neo4j")
-        self._db_query_runner.set_database("neo4j")
+            if self._server_version < ServerVersion(2, 6, 0):
+                raise RuntimeError(
+                    f"AuraDB connection info was provided but GDS version {self._server_version} \
+                        does not support connecting to AuraDB"
+                )
 
-        self._query_runner = AuraDbArrowQueryRunner(
-            gds_query_runner, self._db_query_runner, driver.encrypted, aura_db_connection_info
-        )
+            self._driver_config = gds_query_runner.driver_config()
+            driver = GraphDatabase.driver(
+                aura_db_connection_info.uri, auth=aura_db_connection_info.auth, **self._driver_config
+            )
+            self._db_query_runner: QueryRunner = Neo4jQueryRunner(
+                driver, auto_close=True, bookmarks=bookmarks, server_version=self._server_version
+            )
+
+            # we need to explicitly set these as the default value is None
+            # which signals the driver to use the default configured database
+            # from the dbms.
+            gds_query_runner.set_database("neo4j")
+            self._db_query_runner.set_database("neo4j")
+
+            self._query_runner = AuraDbArrowQueryRunner(
+                gds_query_runner, self._db_query_runner, driver.encrypted, aura_db_connection_info
+            )
+        else:
+            self._query_runner = endpoint
+            self._db_query_runner = endpoint
+            self._server_version = self._query_runner.server_version()
 
         super().__init__(self._query_runner, "gds", self._server_version)
 
@@ -87,6 +94,10 @@ def run_cypher(
         # This will avoid calling valid gds procedures through a raw string
         return self._db_query_runner.run_cypher(query, params, database, False)
 
+    @property
+    def graph(self) -> GraphRemoteProcRunner:
+        return GraphRemoteProcRunner(self._query_runner, f"{self._namespace}.graph", self._server_version)
+
     @property
     def alpha(self) -> AlphaEndpoints:
         return AlphaEndpoints(self._query_runner, "gds.alpha", self._server_version)

graphdatascience/endpoints.py

@@ -3,11 +3,7 @@
     SingleModeAlphaAlgoEndpoints,
 )
 from .call_builder import IndirectAlphaCallBuilder, IndirectBetaCallBuilder
-from .graph.graph_endpoints import (
-    GraphAlphaEndpoints,
-    GraphBetaEndpoints,
-    GraphEndpoints,
-)
+from .graph.graph_endpoints import GraphAlphaEndpoints, GraphBetaEndpoints
 from .model.model_endpoints import (
     ModelAlphaEndpoints,
     ModelBetaEndpoints,
@@ -39,7 +35,6 @@ class DirectEndpoints(
     SingleModeAlgoEndpoints,
     DirectSystemEndpoints,
     DirectUtilEndpoints,
-    GraphEndpoints,
     PipelineEndpoints,
     ModelEndpoints,
     ConfigEndpoints,

graphdatascience/graph/graph_endpoints.py

@@ -1,13 +1,6 @@
 from ..caller_base import CallerBase
 from .graph_alpha_proc_runner import GraphAlphaProcRunner
 from .graph_beta_proc_runner import GraphBetaProcRunner
-from .graph_proc_runner import GraphProcRunner
-
-
-class GraphEndpoints(CallerBase):
-    @property
-    def graph(self) -> GraphProcRunner:
-        return GraphProcRunner(self._query_runner, f"{self._namespace}.graph", self._server_version)
 
 
 class GraphAlphaEndpoints(CallerBase):

graphdatascience/graph/graph_proc_runner.py

@@ -24,7 +24,7 @@
 )
 from .graph_export_runner import GraphExportRunner
 from .graph_object import Graph
-from .graph_project_runner import GraphProjectRunner
+from .graph_project_runner import GraphProjectRemoteRunner, GraphProjectRunner
 from .graph_sample_runner import GraphSampleRunner
 from .graph_type_check import (
     from_graph_type_check,
@@ -41,7 +41,7 @@
 is_neo4j_4_driver = ServerVersion.from_string(neo4j_driver_version) < ServerVersion(5, 0, 0)
 
 
-class GraphProcRunner(UncallableNamespace, IllegalAttrChecker):
+class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
     @staticmethod
     def _path(package: str, resource: str) -> pathlib.Path:
         if sys.version_info >= (3, 9):
@@ -217,11 +217,6 @@ def networkx(self):  # type: ignore
         self._namespace += ".networkx"
         return NXLoader(self._query_runner, self._namespace, self._server_version)
 
-    @property
-    def project(self) -> GraphProjectRunner:
-        self._namespace += ".project"
-        return GraphProjectRunner(self._query_runner, self._namespace, self._server_version)
-
     @property
     @compatible_with("graphProperty", min_inclusive=ServerVersion(2, 5, 0))
     def graphProperty(self) -> GraphPropertyRunner:
@@ -234,11 +229,6 @@ def nodeLabel(self) -> GraphLabelRunner:
         self._namespace += ".nodeLabel"
         return GraphLabelRunner(self._query_runner, self._namespace, self._server_version)
 
-    @property
-    def cypher(self) -> GraphCypherRunner:
-        self._namespace += ".project"
-        return GraphCypherRunner(self._query_runner, self._namespace, self._server_version)
-
     @compatible_with("generate", min_inclusive=ServerVersion(2, 5, 0))
     def generate(self, graph_name: str, node_count: int, average_degree: int, **config: Any) -> GraphCreateResult:
         self._namespace += ".generate"
@@ -568,3 +558,22 @@ def deleteRelationships(self, G: Graph, relationship_type: str) -> "Series[Any]"
             endpoint=self._namespace,
             params=params,
         ).squeeze()
+
+
+class GraphProcRunner(BaseGraphProcRunner):
+    @property
+    def project(self) -> GraphProjectRunner:
+        self._namespace += ".project"
+        return GraphProjectRunner(self._query_runner, self._namespace, self._server_version)
+
+    @property
+    def cypher(self) -> GraphCypherRunner:
+        self._namespace += ".project"
+        return GraphCypherRunner(self._query_runner, self._namespace, self._server_version)
+
+
+class GraphRemoteProcRunner(BaseGraphProcRunner):
+    @property
+    def project(self) -> GraphProjectRemoteRunner:
+        self._namespace += ".project.remoteDb"
+        return GraphProjectRemoteRunner(self._query_runner, self._namespace, self._server_version)

graphdatascience/graph/graph_project_runner.py

@@ -46,10 +46,6 @@ def estimate(self, node_projection: Any, relationship_projection: Any, **config:
     def cypher(self) -> GraphProjectRunner:
         return GraphProjectRunner(self._query_runner, self._namespace + ".cypher", self._server_version)
 
-    @property
-    def remoteDb(self) -> GraphProjectRemoteRunner:
-        return GraphProjectRemoteRunner(self._query_runner, self._namespace + ".remoteDb", self._server_version)
-
 
 class GraphProjectBetaRunner(IllegalAttrChecker):
     @from_graph_type_check
@@ -79,7 +75,7 @@ def subgraph(
 
 
 class GraphProjectRemoteRunner(IllegalAttrChecker):
-    @compatible_with("remoteDb", min_inclusive=ServerVersion(2, 6, 0))
+    @compatible_with("project", min_inclusive=ServerVersion(2, 6, 0))
     def __call__(self, graph_name: str, query: str, **config: Any) -> GraphCreateResult:
         placeholder = "<>"  # host and token will be added by query runner
         params = CallParameters(

graphdatascience/graph_data_science.py

@@ -12,6 +12,7 @@
 from .query_runner.neo4j_query_runner import Neo4jQueryRunner
 from .query_runner.query_runner import QueryRunner
 from .server_version.server_version import ServerVersion
+from graphdatascience.graph.graph_proc_runner import GraphProcRunner
 
 
 class GraphDataScience(DirectEndpoints, UncallableNamespace):
@@ -75,6 +76,10 @@ def __init__(
 
         super().__init__(self._query_runner, "gds", self._server_version)
 
+    @property
+    def graph(self) -> GraphProcRunner:
+        return GraphProcRunner(self._query_runner, f"{self._namespace}.graph", self._server_version)
+
     @property
     def alpha(self) -> AlphaEndpoints:
         return AlphaEndpoints(self._query_runner, "gds.alpha", self._server_version)

graphdatascience/tests/integration/test_remote_graph_ops.py

@@ -37,9 +37,7 @@ def run_around_tests(gds_with_cloud_setup: AuraGraphDataScience) -> Generator[No
 @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 6, 0))
 def test_remote_projection(gds_with_cloud_setup: AuraGraphDataScience) -> None:
-    G, result = gds_with_cloud_setup.graph.project.remoteDb(
-        GRAPH_NAME, "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)"
-    )
+    G, result = gds_with_cloud_setup.graph.project(GRAPH_NAME, "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)")
 
     assert G.name() == GRAPH_NAME
     assert result["nodeCount"] == 3
@@ -48,9 +46,7 @@ def test_remote_projection(gds_with_cloud_setup: AuraGraphDataScience) -> None:
 @pytest.mark.cloud_architecture
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 6, 0))
 def test_remote_write_back(gds_with_cloud_setup: AuraGraphDataScience) -> None:
-    G, result = gds_with_cloud_setup.graph.project.remoteDb(
-        GRAPH_NAME, "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)"
-    )
+    G, result = gds_with_cloud_setup.graph.project(GRAPH_NAME, "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)")
 
     result = gds_with_cloud_setup.pageRank.write(G, writeProperty="score")
 

graphdatascience/tests/unit/conftest.py

@@ -4,8 +4,12 @@
 from pandas import DataFrame
 
 from graphdatascience import QueryRunner
+from graphdatascience.aura_graph_data_science import AuraGraphDataScience
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.graph_data_science import GraphDataScience
+from graphdatascience.query_runner.aura_db_arrow_query_runner import (
+    AuraDbConnectionInfo,
+)
 from graphdatascience.query_runner.cypher_graph_constructor import (
     CypherGraphConstructor,
 )
@@ -109,6 +113,16 @@ def gds(runner: CollectingQueryRunner) -> Generator[GraphDataScience, None, None
     gds.close()
 
 
+@pytest.fixture
+def aura_gds(runner: CollectingQueryRunner) -> Generator[AuraGraphDataScience, None, None]:
+    aura_gds = AuraGraphDataScience(
+        endpoint=runner, auth=("some", "auth"), aura_db_connection_info=AuraDbConnectionInfo("uri", ("some", "auth"))
+    )
+    yield aura_gds
+
+    aura_gds.close()
+
+
 @pytest.fixture(scope="package")
 def server_version() -> ServerVersion:
     return DEFAULT_SERVER_VERSION

graphdatascience/tests/unit/test_graph_ops.py

@@ -2,6 +2,7 @@
 from pandas import DataFrame
 
 from .conftest import CollectingQueryRunner
+from graphdatascience.aura_graph_data_science import AuraGraphDataScience
 from graphdatascience.graph_data_science import GraphDataScience
 from graphdatascience.server_version.server_version import ServerVersion
 
@@ -89,8 +90,8 @@ def test_project_subgraph(runner: CollectingQueryRunner, gds: GraphDataScience)
     }
 
 
-def test_project_remote(runner: CollectingQueryRunner, gds: GraphDataScience) -> None:
-    gds.graph.project.remoteDb("g", "RETURN gds.graph.project.remote(0, 1, null)")
+def test_project_remote(runner: CollectingQueryRunner, aura_gds: AuraGraphDataScience) -> None:
+    aura_gds.graph.project("g", "RETURN gds.graph.project.remote(0, 1, null)")
 
     assert (
         runner.last_query()
@@ -687,9 +688,9 @@ def test_graph_sample_cnarw(runner: CollectingQueryRunner, gds: GraphDataScience
 
 
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 6, 0))
-def test_remote_projection_on_specific_database(runner: CollectingQueryRunner, gds: GraphDataScience) -> None:
-    gds.set_database("bar")
-    G, _ = gds.graph.project.remoteDb("g", "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)")
+def test_remote_projection_on_specific_database(runner: CollectingQueryRunner, aura_gds: AuraGraphDataScience) -> None:
+    aura_gds.set_database("bar")
+    G, _ = aura_gds.graph.project("g", "MATCH (n)-->(m) RETURN gds.graph.project.remote(n, m)")
 
     assert (
         runner.last_query()