Add feature flag to enabled explicit APIs

DarthMax · DarthMax · commit ae2f8cbf7393 · 2025-07-01T17:47:07.000+02:00
diff --git a/graphdatascience/graph_data_science.py b/graphdatascience/graph_data_science.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import warnings
 from types import TracebackType
 from typing import Any, Optional, Type, Union
@@ -10,12 +11,13 @@
 
 from graphdatascience.arrow_client.arrow_authentication import UsernamePasswordAuthentication
 from graphdatascience.arrow_client.arrow_info import ArrowInfo
-from graphdatascience.procedure_surface.cypher.wcc_proc_runner import WccCypherEndpoints
+from graphdatascience.procedure_surface.cypher.wcc_cypher_endpoints import WccCypherEndpoints
 
 from .call_builder import IndirectCallBuilder
 from .endpoints import AlphaEndpoints, BetaEndpoints, DirectEndpoints
 from .error.uncallable_namespace import UncallableNamespace
 from .graph.graph_proc_runner import GraphProcRunner
+from .procedure_surface.api.wcc_endpoints import WccEndpoints
 from .query_runner.arrow_query_runner import ArrowQueryRunner
 from .query_runner.neo4j_query_runner import Neo4jQueryRunner
 from .query_runner.query_runner import QueryRunner
@@ -118,15 +120,20 @@ def __init__(
         self._query_runner.set_show_progress(show_progress)
         super().__init__(self._query_runner, namespace="gds", server_version=self._server_version)
 
-        self._wcc_endpoints = WccCypherEndpoints(self._query_runner)
+        self._wcc_endpoints: Optional[WccEndpoints] = None
+        if os.environ.get("ENABLE_EXPLICIT_ENDPOINTS") is not None:
+            self._wcc_endpoints = WccCypherEndpoints(self._query_runner)
 
     @property
     def graph(self) -> GraphProcRunner:
         return GraphProcRunner(self._query_runner, f"{self._namespace}.graph", self._server_version)
 
-    # @property
-    # def wcc(self) -> WccEndpoints:
-    #     return self._wcc_endpoints
+    @property
+    def wcc(self) -> Union[WccEndpoints, IndirectCallBuilder]:
+        if self._wcc_endpoints is None:
+            return IndirectCallBuilder(self._query_runner, f"gds.{self._namespace}.wcc", self._server_version)
+
+        return self._wcc_endpoints
 
     @property
     def util(self) -> UtilProcRunner:
diff --git a/graphdatascience/procedure_surface/api/wcc_endpoints.py b/graphdatascience/procedure_surface/api/wcc_endpoints.py
@@ -253,6 +253,9 @@ class WccMutateResult:
     node_properties_written: int
     configuration: dict[str, Any]
 
+    def __getitem__(self, item: str) -> Any:
+        return getattr(self, item)
+
 
 @dataclass(frozen=True, repr=True)
 class WccStatsResult:
@@ -263,6 +266,9 @@ class WccStatsResult:
     post_processing_millis: int
     configuration: dict[str, Any]
 
+    def __getitem__(self, item: str) -> Any:
+        return getattr(self, item)
+
 
 @dataclass(frozen=True, repr=True)
 class WccWriteResult:
@@ -274,3 +280,6 @@ class WccWriteResult:
     post_processing_millis: int
     node_properties_written: int
     configuration: dict[str, Any]
+
+    def __getitem__(self, item: str) -> Any:
+        return getattr(self, item)
diff --git a/graphdatascience/session/aura_graph_data_science.py b/graphdatascience/session/aura_graph_data_science.py
@@ -1,12 +1,15 @@
 from __future__ import annotations
 
+import os
 from typing import Any, Callable, Optional, Union
 
 from pandas import DataFrame
 
 from graphdatascience import QueryRunner, ServerVersion
 from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
 from graphdatascience.arrow_client.arrow_info import ArrowInfo
+from graphdatascience.arrow_client.authenticated_arrow_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.write_back_client import WriteBackClient
 from graphdatascience.call_builder import IndirectCallBuilder
 from graphdatascience.endpoints import (
     AlphaRemoteEndpoints,
@@ -15,6 +18,8 @@
 )
 from graphdatascience.error.uncallable_namespace import UncallableNamespace
 from graphdatascience.graph.graph_remote_proc_runner import GraphRemoteProcRunner
+from graphdatascience.procedure_surface.api.wcc_endpoints import WccEndpoints
+from graphdatascience.procedure_surface.arrow.wcc_arrow_endpoints import WccArrowEndpoints
 from graphdatascience.query_runner.arrow_query_runner import ArrowQueryRunner
 from graphdatascience.query_runner.gds_arrow_client import GdsArrowClient
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
@@ -24,15 +29,11 @@
 from graphdatascience.utils.util_remote_proc_runner import UtilRemoteProcRunner
 
 
-class AuraGraphDataScience(DirectEndpoints, UncallableNamespace):
-    """
-    Primary API class for interacting with Neo4j database + Graph Data Science Session.
-    Always bind this object to a variable called `gds`.
-    """
+class AuraGraphDataScienceFactory:
+    """Factory class for creating AuraGraphDataScience instances with all required components."""
 
-    @classmethod
-    def create(
-        cls,
+    def __init__(
+        self,
         session_bolt_connection_info: DbmsConnectionInfo,
         arrow_authentication: Optional[ArrowAuthentication],
         db_endpoint: Optional[Union[Neo4jQueryRunner, DbmsConnectionInfo]],
@@ -41,74 +42,114 @@ def create(
         arrow_tls_root_certs: Optional[bytes] = None,
         bookmarks: Optional[Any] = None,
         show_progress: bool = True,
-    ) -> AuraGraphDataScience:
-        session_bolt_query_runner = Neo4jQueryRunner.create_for_session(
-            endpoint=session_bolt_connection_info.uri,
-            auth=session_bolt_connection_info.get_auth(),
-            show_progress=show_progress,
+    ):
+        self.session_bolt_connection_info = session_bolt_connection_info
+        self.arrow_authentication = arrow_authentication
+        self.db_endpoint = db_endpoint
+        self.delete_fn = delete_fn
+        self.arrow_disable_server_verification = arrow_disable_server_verification
+        self.arrow_tls_root_certs = arrow_tls_root_certs
+        self.bookmarks = bookmarks
+        self.show_progress = show_progress
+
+    def create(self) -> AuraGraphDataScience:
+        """Create and configure an AuraGraphDataScience instance."""
+        session_bolt_query_runner = self._create_session_bolt_query_runner()
+        arrow_info = ArrowInfo.create(session_bolt_query_runner)
+        session_arrow_query_runner = self._create_session_arrow_query_runner(session_bolt_query_runner, arrow_info)
+        session_arrow_client = self._create_session_arrow_client(arrow_info, session_bolt_query_runner)
+        gds_version = session_bolt_query_runner.server_version()
+
+        session_query_runner: QueryRunner
+
+        if self.db_endpoint is not None:
+            db_bolt_query_runner = self._create_db_bolt_query_runner()
+            session_query_runner = SessionQueryRunner.create(
+                session_arrow_query_runner, db_bolt_query_runner, session_arrow_client, self.show_progress
+            )
+            wcc_endpoints = self._create_wcc_endpoints(arrow_info, session_bolt_query_runner, db_bolt_query_runner)
+        else:
+            session_query_runner = StandaloneSessionQueryRunner(session_arrow_query_runner)
+            wcc_endpoints = self._create_wcc_endpoints(arrow_info, session_bolt_query_runner, None)
+
+        return AuraGraphDataScience(
+            query_runner=session_query_runner,
+            wcc_endpoints=wcc_endpoints,
+            delete_fn=self.delete_fn,
+            gds_version=gds_version,
         )
 
-        arrow_info = ArrowInfo.create(session_bolt_query_runner)
-        session_arrow_query_runner = ArrowQueryRunner.create(
+    def _create_session_bolt_query_runner(self) -> Neo4jQueryRunner:
+        return Neo4jQueryRunner.create_for_session(
+            endpoint=self.session_bolt_connection_info.uri,
+            auth=self.session_bolt_connection_info.get_auth(),
+            show_progress=self.show_progress,
+        )
+
+    def _create_session_arrow_query_runner(
+        self, session_bolt_query_runner: Neo4jQueryRunner, arrow_info: ArrowInfo
+    ) -> ArrowQueryRunner:
+        return ArrowQueryRunner.create(
             fallback_query_runner=session_bolt_query_runner,
             arrow_info=arrow_info,
-            arrow_authentication=arrow_authentication,
+            arrow_authentication=self.arrow_authentication,
             encrypted=session_bolt_query_runner.encrypted(),
-            disable_server_verification=arrow_disable_server_verification,
-            tls_root_certs=arrow_tls_root_certs,
+            disable_server_verification=self.arrow_disable_server_verification,
+            tls_root_certs=self.arrow_tls_root_certs,
         )
 
-        # TODO: merge with the gds_arrow_client created inside ArrowQueryRunner
-        session_arrow_client = GdsArrowClient.create(
+    def _create_session_arrow_client(
+        self, arrow_info: ArrowInfo, session_bolt_query_runner: Neo4jQueryRunner
+    ) -> GdsArrowClient:
+        return GdsArrowClient.create(
             arrow_info,
-            arrow_authentication,
+            self.arrow_authentication,
             session_bolt_query_runner.encrypted(),
-            arrow_disable_server_verification,
-            arrow_tls_root_certs,
+            self.arrow_disable_server_verification,
+            self.arrow_tls_root_certs,
         )
 
-        gds_version = session_bolt_query_runner.server_version()
-
-        if db_endpoint is not None:
-            if isinstance(db_endpoint, Neo4jQueryRunner):
-                db_bolt_query_runner = db_endpoint
-            else:
-                db_bolt_query_runner = Neo4jQueryRunner.create_for_db(
-                    db_endpoint.uri,
-                    db_endpoint.get_auth(),
-                    aura_ds=True,
-                    show_progress=False,
-                    database=db_endpoint.database,
-                )
-            db_bolt_query_runner.set_bookmarks(bookmarks)
-
-            session_query_runner = SessionQueryRunner.create(
-                session_arrow_query_runner, db_bolt_query_runner, session_arrow_client, show_progress
-            )
-            return cls(
-                query_runner=session_query_runner,
-                delete_fn=delete_fn,
-                gds_version=gds_version,
+    def _create_db_bolt_query_runner(self) -> Neo4jQueryRunner:
+        if isinstance(self.db_endpoint, Neo4jQueryRunner):
+            db_bolt_query_runner = self.db_endpoint
+        elif isinstance(self.db_endpoint, DbmsConnectionInfo):
+            db_bolt_query_runner = Neo4jQueryRunner.create_for_db(
+                self.db_endpoint.uri,
+                self.db_endpoint.get_auth(),
+                aura_ds=True,
+                show_progress=False,
+                database=self.db_endpoint.database,
             )
         else:
-            standalone_query_runner = StandaloneSessionQueryRunner(session_arrow_query_runner)
-            return cls(
-                query_runner=standalone_query_runner,
-                delete_fn=delete_fn,
-                gds_version=gds_version,
+            raise ValueError("db_endpoint must be a Neo4jQueryRunner or a DbmsConnectionInfo")
+
+        db_bolt_query_runner.set_bookmarks(self.bookmarks)
+        return db_bolt_query_runner
+
+    def _create_wcc_endpoints(
+        self, arrow_info: ArrowInfo, session_bolt_query_runner: Neo4jQueryRunner, db_query_runner: Optional[QueryRunner]
+    ) -> Optional[WccEndpoints]:
+        wcc_endpoints: Optional[WccEndpoints] = None
+        if os.environ.get("ENABLE_EXPLICIT_ENDPOINTS") is not None:
+            arrow_client = AuthenticatedArrowClient.create(
+                arrow_info,
+                self.arrow_authentication,
+                session_bolt_query_runner.encrypted(),
+                self.arrow_disable_server_verification,
+                self.arrow_tls_root_certs,
             )
 
-    def __init__(
-        self,
-        query_runner: QueryRunner,
-        delete_fn: Callable[[], bool],
-        gds_version: ServerVersion,
-    ):
-        self._query_runner = query_runner
-        self._delete_fn = delete_fn
-        self._server_version = gds_version
+            write_back_client = WriteBackClient(arrow_client, db_query_runner) if db_query_runner is not None else None
 
-        super().__init__(self._query_runner, namespace="gds", server_version=self._server_version)
+            wcc_endpoints = WccArrowEndpoints(arrow_client, write_back_client)
+        return wcc_endpoints
+
+
+class AuraGraphDataScience(DirectEndpoints, UncallableNamespace):
+    """
+    Primary API class for interacting with Neo4j database + Graph Data Science Session.
+    Always bind this object to a variable called `gds`.
+    """
 
     def run_cypher(
         self,
@@ -133,6 +174,20 @@ def run_cypher(
         """
         return self._query_runner.run_cypher(query, params, database, False)
 
+    def __init__(
+        self,
+        query_runner: QueryRunner,
+        delete_fn: Callable[[], bool],
+        gds_version: ServerVersion,
+        wcc_endpoints: Optional[WccEndpoints] = None,
+    ):
+        self._query_runner = query_runner
+        self._delete_fn = delete_fn
+        self._server_version = gds_version
+        self._wcc_endpoints = wcc_endpoints
+
+        super().__init__(self._query_runner, namespace="gds", server_version=self._server_version)
+
     @property
     def graph(self) -> GraphRemoteProcRunner:
         return GraphRemoteProcRunner(self._query_runner, f"{self._namespace}.graph", self._server_version)
@@ -149,6 +204,13 @@ def alpha(self) -> AlphaRemoteEndpoints:
     def beta(self) -> BetaEndpoints:
         return BetaEndpoints(self._query_runner, "gds.beta", self._server_version)
 
+    @property
+    def wcc(self) -> Union[WccEndpoints, IndirectCallBuilder]:
+        if self._wcc_endpoints is None:
+            return IndirectCallBuilder(self._query_runner, f"gds.{self._namespace}.wcc", self._server_version)
+
+        return self._wcc_endpoints
+
     def __getattr__(self, attr: str) -> IndirectCallBuilder:
         return IndirectCallBuilder(self._query_runner, f"gds.{attr}", self._server_version)
 
diff --git a/graphdatascience/session/dedicated_sessions.py b/graphdatascience/session/dedicated_sessions.py
@@ -11,7 +11,7 @@
 from graphdatascience.session.aura_api import AuraApi
 from graphdatascience.session.aura_api_responses import SessionDetails
 from graphdatascience.session.aura_api_token_authentication import AuraApiTokenAuthentication
-from graphdatascience.session.aura_graph_data_science import AuraGraphDataScience
+from graphdatascience.session.aura_graph_data_science import AuraGraphDataScience, AuraGraphDataScienceFactory
 from graphdatascience.session.cloud_location import CloudLocation
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
 from graphdatascience.session.session_info import SessionInfo
@@ -210,9 +210,9 @@ def _construct_client(
         arrow_authentication: ArrowAuthentication,
         db_runner: Optional[Neo4jQueryRunner],
     ) -> AuraGraphDataScience:
-        return AuraGraphDataScience.create(
+        return AuraGraphDataScienceFactory(
             session_bolt_connection_info=session_bolt_connection_info,
             arrow_authentication=arrow_authentication,
             db_endpoint=db_runner,
             delete_fn=lambda: self._aura_api.delete_session(session_id=session_id),
-        )
+        ).create()
diff --git a/graphdatascience/tests/integration/conftest.py b/graphdatascience/tests/integration/conftest.py
@@ -10,7 +10,7 @@
 from graphdatascience.graph_data_science import GraphDataScience
 from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 from graphdatascience.server_version.server_version import ServerVersion
-from graphdatascience.session.aura_graph_data_science import AuraGraphDataScience
+from graphdatascience.session.aura_graph_data_science import AuraGraphDataScience, AuraGraphDataScienceFactory
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
 
 URI = os.environ.get("NEO4J_URI", "bolt://localhost:7687")
@@ -92,12 +92,13 @@ def gds_without_arrow() -> Generator[GraphDataScience, None, None]:
 
 @pytest.fixture(scope="package", autouse=False)
 def gds_with_cloud_setup(request: pytest.FixtureRequest) -> Generator[AuraGraphDataScience, None, None]:
-    _gds = AuraGraphDataScience.create(
+    _gds = AuraGraphDataScienceFactory(
         session_bolt_connection_info=DbmsConnectionInfo(URI, AUTH[0], AUTH[1]),
         arrow_authentication=UsernamePasswordAuthentication(AUTH[0], AUTH[1]),
         db_endpoint=DbmsConnectionInfo(AURA_DB_URI, AURA_DB_AUTH[0], AURA_DB_AUTH[1]),
         delete_fn=lambda: True,
-    )
+    ).create()
+
     _gds.set_database(DB)
 
     yield _gds
@@ -107,12 +108,12 @@ def gds_with_cloud_setup(request: pytest.FixtureRequest) -> Generator[AuraGraphD
 
 @pytest.fixture(scope="package", autouse=False)
 def standalone_aura_gds() -> Generator[AuraGraphDataScience, None, None]:
-    _gds = AuraGraphDataScience.create(
+    _gds = AuraGraphDataScienceFactory(
         session_bolt_connection_info=DbmsConnectionInfo(URI, AUTH[0], AUTH[1]),
         arrow_authentication=UsernamePasswordAuthentication(AUTH[0], AUTH[1]),
         db_endpoint=None,
         delete_fn=lambda: True,
-    )
+    ).create()
 
     yield _gds
 
diff --git a/graphdatascience/tests/integration/test_remote_graph_ops.py b/graphdatascience/tests/integration/test_remote_graph_ops.py
@@ -62,7 +62,7 @@ def test_remote_projection_and_writeback_custom_database_name(gds_with_cloud_set
         assert projection_result["nodeCount"] == 2
         assert projection_result["relationshipCount"] == 1
 
-        write_result = gds_with_cloud_setup.wcc.write(G, writeProperty="wcc")
+        write_result = gds_with_cloud_setup.wcc.write(G, writeProperty="wcc")  # type: ignore
 
         assert write_result["nodePropertiesWritten"] == 2
         count_wcc_nodes_query = "MATCH (n WHERE n.wcc IS NOT NULL) RETURN count(*) AS c"
@@ -234,6 +234,6 @@ def test_empty_graph_write_back(
 
     assert G.node_count() == 0
 
-    result = gds_with_cloud_setup.wcc.write(G, writeProperty="wcc")
+    result = gds_with_cloud_setup.wcc.write(G, writeProperty="wcc")  # type: ignore
 
     assert result["nodePropertiesWritten"] == 0
diff --git a/graphdatascience/tests/integration/test_simple_algo.py b/graphdatascience/tests/integration/test_simple_algo.py
diff --git a/graphdatascience/tests/unit/conftest.py b/graphdatascience/tests/unit/conftest.py
