Expose v2 endpoints on plugin

using the cypher-api here

Author: FlorentinD

graphdatascience/__init__.py

@@ -10,6 +10,7 @@
 from .pipeline.lp_training_pipeline import LPTrainingPipeline
 from .pipeline.nc_training_pipeline import NCTrainingPipeline
 from .pipeline.nr_training_pipeline import NRTrainingPipeline
+from .plugin_v2_endpoints import PluginV2Endpoints
 from .query_runner.query_runner import QueryRunner
 from .server_version.server_version import ServerVersion
 from .session.gds_sessions import GdsSessions
@@ -33,4 +34,5 @@
     "NRModel",
     "GraphSageModel",
     "SimpleRelEmbeddingModel",
+    "PluginV2Endpoints",
 ]

graphdatascience/graph_data_science.py

@@ -8,6 +8,7 @@
 from neo4j import Driver
 from pandas import DataFrame
 
+from graphdatascience.plugin_v2_endpoints import PluginV2Endpoints
 from graphdatascience.query_runner.arrow_authentication import UsernamePasswordAuthentication
 from graphdatascience.query_runner.query_mode import QueryMode
 from graphdatascience.topological_lp.topological_lp_runner import TopologicalLPRunner
@@ -87,15 +88,17 @@ def __init__(
         if aura_ds:
             GraphDataScience._validate_endpoint(endpoint)
 
+        neo4j_query_runner: Neo4jQueryRunner | None = None
         if isinstance(endpoint, QueryRunner):
             self._query_runner = endpoint
         else:
             db_auth = None
             if auth:
                 db_auth = neo4j.basic_auth(*auth)
-            self._query_runner = Neo4jQueryRunner.create_for_db(
+            neo4j_query_runner = Neo4jQueryRunner.create_for_db(
                 endpoint, db_auth, aura_ds, database, bookmarks, show_progress
             )
+            self._query_runner = neo4j_query_runner
 
         self._server_version = self._query_runner.server_version()
 
@@ -130,6 +133,14 @@ def __init__(
                 connection_string_override=None if arrow is True else arrow,
             )
 
+        arrow_client = (
+            None if not isinstance(self._query_runner, ArrowQueryRunner) else self._query_runner._gds_arrow_client
+        )
+        if neo4j_query_runner:
+            self._v2_endpoints: PluginV2Endpoints | None = PluginV2Endpoints(neo4j_query_runner, arrow_client)
+        else:
+            self._v2_endpoints = None
+
         self._query_runner.set_show_progress(show_progress)
         super().__init__(self._query_runner, namespace="gds", server_version=self._server_version)
 
@@ -265,6 +276,17 @@ def driver_config(self) -> dict[str, Any]:
         """
         return self._query_runner.driver_config()
 
+    @property
+    def v2(self) -> PluginV2Endpoints:
+        """
+        Return preview v2 endpoints. These endpoints may change without warning.
+        These endpoints are a preview of the API for the next major version of this library.
+        """
+        if not self._v2_endpoints:
+            raise RuntimeError("v2 endpoints are not available.")
+
+        return self._v2_endpoints
+
     @classmethod
     def from_neo4j_driver(
         cls: Type[GraphDataScience],

graphdatascience/plugin_v2_endpoints.py

@@ -0,0 +1,7 @@
+from graphdatascience.procedure_surface.cypher.catalog_cypher_endpoints import (
+    CatalogCypherEndpoints,
+    GraphProjectResult,
+    GraphWithProjectResult,
+)
+
+__all__ = ["CatalogCypherEndpoints", "GraphWithProjectResult", "GraphProjectResult"]

graphdatascience/procedure_surface/cypher/__init__.py

@@ -16,9 +16,10 @@
 from graphdatascience.procedure_surface.api.catalog.graph_info import GraphInfo, GraphInfoWithDegrees
 from graphdatascience.procedure_surface.api.catalog.graph_sampling_endpoints import GraphSamplingEndpoints
 from graphdatascience.procedure_surface.cypher.catalog.graph_backend_cypher import get_graph
+from graphdatascience.query_runner.gds_arrow_client import GdsArrowClient
+from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 
 from ...call_parameters import CallParameters
-from ...query_runner.query_runner import QueryRunner
 from ..api.base_result import BaseResult
 from ..utils.config_converter import ConfigConverter
 from .catalog.graph_sampling_cypher_endpoints import GraphSamplingCypherEndpoints
@@ -28,14 +29,15 @@
 
 
 class CatalogCypherEndpoints(CatalogEndpoints):
-    def __init__(self, query_runner: QueryRunner):
-        self._query_runner = query_runner
+    def __init__(self, cypher_runner: Neo4jQueryRunner, arrow_client: GdsArrowClient | None = None):
+        self.cypher_runner = cypher_runner
+        self._arrow_client = arrow_client
 
     def list(self, G: GraphV2 | str | None = None) -> list[GraphInfoWithDegrees]:
         graph_name = G if isinstance(G, str) else G.name() if G is not None else None
         params = CallParameters(graphName=graph_name) if graph_name else CallParameters()
 
-        result = self._query_runner.call_procedure(endpoint="gds.graph.list", params=params)
+        result = self.cypher_runner.call_procedure(endpoint="gds.graph.list", params=params)
         return [GraphInfoWithDegrees(**row.to_dict()) for _, row in result.iterrows()]
 
     def drop(self, G: GraphV2 | str, fail_if_missing: bool = True) -> GraphInfo | None:
@@ -47,7 +49,7 @@ def drop(self, G: GraphV2 | str, fail_if_missing: bool = True) -> GraphInfo | No
             else CallParameters(graphName=graph_name)
         )
 
-        result = self._query_runner.call_procedure(endpoint="gds.graph.drop", params=params)
+        result = self.cypher_runner.call_procedure(endpoint="gds.graph.drop", params=params)
         if len(result) > 0:
             return GraphInfo(**result.iloc[0].to_dict())
         else:
@@ -64,6 +66,7 @@ def project(
         job_id: str | None = None,
         sudo: bool = False,
         username: str | None = None,
+        log_progress: bool = True,
     ) -> GraphWithProjectResult:
         config = ConfigConverter.convert_to_gds_config(
             nodeProperties=node_properties,
@@ -82,9 +85,11 @@ def project(
         )
         params.ensure_job_id_in_config()
 
-        result = self._query_runner.call_procedure(endpoint="gds.graph.project", params=params).squeeze()
+        result = self.cypher_runner.call_procedure(
+            endpoint="gds.graph.project", params=params, logging=log_progress
+        ).squeeze()
         project_result = GraphProjectResult(**result.to_dict())
-        return GraphWithProjectResult(get_graph(project_result.graph_name, self._query_runner), project_result)
+        return GraphWithProjectResult(get_graph(project_result.graph_name, self.cypher_runner), project_result)
 
     def filter(
         self,
@@ -94,6 +99,7 @@ def filter(
         relationship_filter: str,
         concurrency: int | None = None,
         job_id: str | None = None,
+        log_progress: bool = True,
     ) -> GraphWithFilterResult:
         config = ConfigConverter.convert_to_gds_config(
             concurrency=concurrency,
@@ -109,8 +115,10 @@ def filter(
         )
         params.ensure_job_id_in_config()
 
-        result = self._query_runner.call_procedure(endpoint="gds.graph.filter", params=params).squeeze()
-        return GraphWithFilterResult(get_graph(graph_name, self._query_runner), GraphFilterResult(**result.to_dict()))
+        result = self.cypher_runner.call_procedure(
+            endpoint="gds.graph.filter", params=params, logging=log_progress
+        ).squeeze()
+        return GraphWithFilterResult(get_graph(graph_name, self.cypher_runner), GraphFilterResult(**result.to_dict()))
 
     def generate(
         self,
@@ -151,26 +159,28 @@ def generate(
 
         params.ensure_job_id_in_config()
 
-        result = self._query_runner.call_procedure(endpoint="gds.graph.generate", params=params).squeeze()
+        result = self.cypher_runner.call_procedure(
+            endpoint="gds.graph.generate", params=params, logging=log_progress
+        ).squeeze()
         return GraphWithGenerationStats(
-            get_graph(graph_name, self._query_runner), GraphGenerationStats(**result.to_dict())
+            get_graph(graph_name, self.cypher_runner), GraphGenerationStats(**result.to_dict())
         )
 
     @property
     def sample(self) -> GraphSamplingEndpoints:
-        return GraphSamplingCypherEndpoints(self._query_runner)
+        return GraphSamplingCypherEndpoints(self.cypher_runner)
 
     @property
     def node_labels(self) -> NodeLabelCypherEndpoints:
-        return NodeLabelCypherEndpoints(self._query_runner)
+        return NodeLabelCypherEndpoints(self.cypher_runner)
 
     @property
     def node_properties(self) -> NodePropertiesCypherEndpoints:
-        return NodePropertiesCypherEndpoints(self._query_runner)
+        return NodePropertiesCypherEndpoints(self.cypher_runner, self._arrow_client)
 
     @property
     def relationships(self) -> RelationshipCypherEndpoints:
-        return RelationshipCypherEndpoints(self._query_runner)
+        return RelationshipCypherEndpoints(self.cypher_runner, self._arrow_client)
 
 
 class GraphProjectResult(BaseResult):

graphdatascience/procedure_surface/cypher/catalog_cypher_endpoints.py

@@ -2,8 +2,8 @@
 
 from pandas import DataFrame
 
-from graphdatascience import QueryRunner
 from graphdatascience.query_runner.query_mode import QueryMode
+from graphdatascience.query_runner.query_runner import QueryRunner
 
 
 def transpose_property_columns(result: DataFrame, list_node_labels: bool) -> DataFrame:

graphdatascience/procedure_surface/utils/result_utils.py

@@ -152,10 +152,16 @@ def graph(self) -> CatalogArrowEndpoints:
 
     @property
     def config(self) -> ConfigArrowEndpoints:
+        """
+        Return configuration-related endpoints.
+        """
         return ConfigArrowEndpoints(self._arrow_client)
 
     @property
     def system(self) -> SystemArrowEndpoints:
+        """
+        Return system-related endpoints.
+        """
         return SystemArrowEndpoints(self._arrow_client)
 
     ## Algorithms

graphdatascience/session/session_v2_endpoints.py

@@ -1,22 +1,29 @@
 import logging
 import os
+import subprocess
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
 from typing import Generator
 
+import dotenv
 import pytest
 from dateutil.relativedelta import relativedelta
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.network import Network
 from testcontainers.core.waiting_utils import wait_for_logs
+from testcontainers.neo4j import Neo4jContainer
 
 from graphdatascience.arrow_client.arrow_authentication import UsernamePasswordAuthentication
 from graphdatascience.arrow_client.arrow_info import ArrowInfo
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
 from graphdatascience.tests.integrationV2.conftest import inside_ci
+from graphdatascience.tests.integrationV2.procedure_surface.gds_api_spec import (
+    EndpointWithModesSpec,
+    resolve_spec_from_file,
+)
 
 LOGGER = logging.getLogger(__name__)
 
@@ -49,6 +56,49 @@ def network() -> Generator[Network, None, None]:
         yield network
 
 
+@pytest.fixture(scope="session")
+def gds_api_spec(tmp_path_factory: pytest.TempPathFactory) -> Generator[list[EndpointWithModesSpec], None, None]:
+    provided_spec_file = os.environ.get("GDS_API_SPEC_FILE")
+
+    spec_file: Path | None = None
+    if provided_spec_file:
+        spec_file = Path(provided_spec_file)
+
+    if spec_file and not spec_file.exists():
+        raise FileNotFoundError(f"GDS_API_SPEC_FILE is set to '{spec_file}', but the file does not exist.")
+
+    if not spec_file:
+        spec_dir = tmp_path_factory.mktemp("gds_api_spec")
+        spec_file = spec_dir / "gds-api-spec.json"
+
+        # allow for caching
+        if not spec_file.exists():
+            download_gds_api_spec(spec_file)
+
+    # Adjust the path to pull from graph-analytics
+    yield resolve_spec_from_file(spec_file)
+
+
+def download_gds_api_spec(destination: Path) -> None:
+    import requests
+
+    url = "https://raw.githubusercontent.com/neo-technology/graph-analytics/refs/heads/master/tools/gds-api-spec/gds-api-spec.json"
+    gh_token = os.environ.get("GITHUB_TOKEN")
+    if not gh_token:
+        try:
+            result = subprocess.run(["gh", "auth", "token"], capture_output=True, text=True, check=True)
+            gh_token = result.stdout.strip()
+        except (subprocess.CalledProcessError, FileNotFoundError) as e:
+            raise ValueError("Failed to get GitHub token. Set GITHUB_TOKEN or authenticate with gh CLI.") from e
+
+    headers = {"Authorization": f"Token {gh_token}"}
+    response = requests.get(url, headers=headers)
+    response.raise_for_status()
+
+    with open(destination, "wb") as f:
+        f.write(response.content)
+
+
 def latest_neo4j_version() -> str:
     today = datetime.now()
     previous_month = today - relativedelta(months=1)
@@ -156,6 +206,60 @@ def start_database(logs_dir: Path, network: Network) -> Generator[DbmsConnection
                 f.write(stdout.decode("utf-8"))
 
 
+def start_gds_plugin_database(
+    logs_dir: Path, tmp_path_factory: pytest.TempPathFactory
+) -> Generator[Neo4jContainer, None, None]:
+    neo4j_image = os.getenv("NEO4J_DATABASE_IMAGE", "neo4j:enterprise")
+
+    dotenv.load_dotenv("graphdatascience/tests/test.env", override=True)
+    GDS_LICENSE_KEY = os.getenv("GDS_LICENSE_KEY")
+
+    db_logs_dir = logs_dir / "cypher_surface" / "db_logs"
+    db_logs_dir.mkdir(parents=True)
+    db_logs_dir.chmod(0o777)
+
+    neo4j_container = (
+        Neo4jContainer(
+            image=neo4j_image,
+        )
+        .with_env("NEO4J_ACCEPT_LICENSE_AGREEMENT", "yes")
+        .with_env("NEO4J_PLUGINS", '["graph-data-science"]')
+        .with_env("NEO4J_gds_arrow_enabled", "true")
+        .with_env("NEO4J_gds_arrow_listen__address", "0.0.0.0:8491")
+        .with_exposed_ports(8491)
+        .with_volume_mapping(db_logs_dir, "/logs", mode="rw")
+    )
+
+    if GDS_LICENSE_KEY is not None:
+        license_dir = tmp_path_factory.mktemp("gds_license")
+        license_dir.chmod(0o755)
+        license_file = os.path.join(license_dir, "license_key")
+        with open(license_file, "w") as f:
+            f.write(GDS_LICENSE_KEY)
+
+        neo4j_container.with_volume_mapping(
+            license_dir,
+            "/licenses",
+        )
+        neo4j_container.with_env("NEO4J_gds_enterprise_license__file", "/licenses/license_key")
+
+    with neo4j_container as neo4j_db:
+        try:
+            wait_for_logs(neo4j_db, "Started.")
+            yield neo4j_db
+        finally
+            stdout, stderr = neo4j_db.get_logs()
+            if stderr:
+                print(f"Error logs from Neo4j container:\n{stderr}")
+
+            if inside_ci():
+                print(f"Neo4j container logs:\n{stdout}")
+
+            out_file = db_logs_dir / "stdout.log"
+            with open(out_file, "w") as f:
+                f.write(stdout.decode("utf-8"))
+
+
 def create_db_query_runner(neo4j_connection: DbmsConnectionInfo) -> Generator[Neo4jQueryRunner, None, None]:
     query_runner = Neo4jQueryRunner.create_for_db(
         f"bolt://{neo4j_connection.uri}",

graphdatascience/tests/integrationV2/procedure_surface/conftest.py

@@ -9,6 +9,7 @@
 from graphdatascience.procedure_surface.cypher.catalog.relationship_cypher_endpoints import (
     RelationshipCypherEndpoints,
 )
+from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 from graphdatascience.tests.integrationV2.procedure_surface.cypher.cypher_graph_helper import create_graph
 
 
@@ -45,7 +46,7 @@ def sample_graph(query_runner: QueryRunner) -> Generator[GraphV2, None, None]:
 
 @pytest.fixture
 def relationship_endpoints(
-    query_runner: QueryRunner,
+    query_runner: Neo4jQueryRunner,
 ) -> Generator[RelationshipCypherEndpoints, None, None]:
     yield RelationshipCypherEndpoints(query_runner)