write integration tests for v2 GdsArrowClient

Author: DarthMax

graphdatascience/arrow_client/v2/gds_arrow_client.py

@@ -82,12 +82,11 @@ def get_node_properties(
 
         return JobClient.run_job(self._flight_client, "v2/graph.nodeProperties.stream", config)
 
-    def get_node_labels(
+    def get_nodes(
         self,
         graph_name: str,
-        node_label: str,
         *,
-        node_filter: str,
+        node_filter: str | None = None,
         log_progress: bool = True,
         concurrency: int | None = None,
         job_id: str | None = None,
@@ -99,8 +98,6 @@ def get_node_labels(
         ----------
         graph_name
            The name of the graph
-        node_label
-            The node label to stream back.
         node_filter
             A Cypher predicate for filtering nodes in the input graph.
         log_progress
@@ -117,14 +114,14 @@ def get_node_labels(
         """
         config = ConfigConverter.convert_to_gds_config(
             graph_name=graph_name,
-            node_label=node_label,
+            node_label="__IGNORED__",
             node_filter=node_filter,
             concurrency=concurrency,
             log_progress=log_progress,
             job_id=job_id,
         )
 
-        return JobClient.run_job(self._flight_client, "v2/graph.nodeLabel.stream", config)
+        return JobClient.run_job_and_wait(self._flight_client, "v2/graph.nodeLabel.stream", config, log_progress)
 
     def get_relationships(
         self,
@@ -325,15 +322,15 @@ def triplet_load_done(self, job_id: str) -> None:
         """
         self._flight_client.do_action_with_retry("v2/graph.project.fromTriplets.done", {"jobId": job_id})
 
-    def upload_data(
+    def upload_nodes(
         self,
         job_id: str,
         data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
         batch_size: int = 10000,
         progress_callback: Callable[[int], None] = lambda x: None,
     ) -> None:
         """
-        Uploads data to the server for a given job.
+        Uploads node data to the server for a given job.
 
         Parameters
         ----------
@@ -346,7 +343,62 @@ def upload_data(
         progress_callback
             A callback function that is called with the number of rows uploaded after each batch
         """
+        self._upload_data("graph.project.fromTables.nodes", job_id, data, batch_size, progress_callback)
 
+    def upload_relationships(
+        self,
+        job_id: str,
+        data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
+        batch_size: int = 10000,
+        progress_callback: Callable[[int], None] = lambda x: None,
+    ) -> None:
+        """
+        Uploads relationship data to the server for a given job.
+
+        Parameters
+        ----------
+        job_id
+            The job id of the import process
+        data
+            The data to upload
+        batch_size
+            The number of rows per batch
+        progress_callback
+            A callback function that is called with the number of rows uploaded after each batch
+        """
+        self._upload_data("graph.project.fromTables.relationships", job_id, data, batch_size, progress_callback)
+
+    def upload_triplets(
+        self,
+        job_id: str,
+        data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
+        batch_size: int = 10000,
+        progress_callback: Callable[[int], None] = lambda x: None,
+    ) -> None:
+        """
+        Uploads triplet data to the server for a given job.
+
+        Parameters
+        ----------
+        job_id
+            The job id of the import process
+        data
+            The data to upload
+        batch_size
+            The number of rows per batch
+        progress_callback
+            A callback function that is called with the number of rows uploaded after each batch
+        """
+        self._upload_data("graph.project.fromTriplets", job_id, data, batch_size, progress_callback)
+
+    def _upload_data(
+        self,
+        endpoint: str,
+        job_id: str,
+        data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
+        batch_size: int = 10000,
+        progress_callback: Callable[[int], None] = lambda x: None,
+    ) -> None:
         match data:
             case pyarrow.Table():
                 batches = data.to_batches(batch_size)
@@ -356,9 +408,11 @@ def upload_data(
                 batches = data
 
         flight_descriptor = {
-            "name": job_id,
+            "name": endpoint,
             "version": ArrowEndpointVersion.V2.version(),
-            "body": {},
+            "body": {
+                "jobId": job_id,
+            },
         }
         upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
 

graphdatascience/tests/integrationV2/arrow_client/v2/test_gds_arrow_client.py

@@ -0,0 +1,134 @@
+from pathlib import Path
+from typing import Generator
+
+import numpy as np
+import pandas as pd
+import pytest
+from testcontainers.core.network import Network
+
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.gds_arrow_client import GdsArrowClient
+from graphdatascience.procedure_surface.api.catalog import GraphV2
+from graphdatascience.procedure_surface.arrow.catalog import CatalogArrowEndpoints
+from graphdatascience.tests.integrationV2.conftest import GdsSessionConnectionInfo, create_arrow_client, start_session
+from graphdatascience.tests.integrationV2.procedure_surface.arrow.graph_creation_helper import create_graph
+
+
+@pytest.fixture(scope="package")
+def session_connection(
+    network: Network, password_dir: Path, logs_dir: Path
+) -> Generator[GdsSessionConnectionInfo, None, None]:
+    yield from start_session(logs_dir, network, password_dir)
+
+
+@pytest.fixture(scope="package")
+def arrow_client(session_connection: GdsSessionConnectionInfo) -> AuthenticatedArrowClient:
+    return create_arrow_client(session_connection)
+
+
+@pytest.fixture(scope="package")
+def gds_arrow_client(arrow_client: AuthenticatedArrowClient) -> GdsArrowClient:
+    return GdsArrowClient(arrow_client)
+
+
+@pytest.fixture
+def sample_graph(arrow_client: AuthenticatedArrowClient) -> Generator[GraphV2, None, None]:
+    gdl = """
+        CREATE
+        (a: Node:Foo {prop1: 1, prop2: 42.0}),
+        (b: Node {prop1: 2, prop2: 43.0}),
+        (c: Node:Foo {prop1: 3, prop2: 44.0}),
+        
+        (a)-[:REL {relX: 1, relY: 42}]->(b),
+        (b)-[:REL {relX: 2, relY: 43}]->(c),
+        (c)-[:REL2 {relX: 1, relY: 2}]->(a),
+    """
+
+    with create_graph(arrow_client, "g", gdl) as G:
+        yield G
+
+
+def test_stream_node_label(gds_arrow_client: GdsArrowClient, sample_graph: GraphV2) -> None:
+    job_id = gds_arrow_client.get_nodes(sample_graph.name(), node_filter="n.prop1 > 1")
+    result = gds_arrow_client.stream(sample_graph.name(), job_id)
+
+    assert ["nodeId"] == list(result.columns)
+    assert len(result) == 2
+
+
+def test_stream_node_properties(gds_arrow_client: GdsArrowClient, sample_graph: GraphV2) -> None:
+    job_id = gds_arrow_client.get_node_properties(sample_graph.name(), node_properties=["prop1", "prop2"])
+    result = gds_arrow_client.stream(sample_graph.name(), job_id)
+
+    assert len(result) == 3
+    assert "nodeId" in result.columns
+    assert "prop1" in result.columns
+    assert "prop2" in result.columns
+    assert set(result["prop1"].tolist()) == {1, 2, 3}
+    assert set(result["prop2"].tolist()) == {42.0, 43.0, 44.0}
+
+
+def test_stream_relationship_properties(gds_arrow_client: GdsArrowClient, sample_graph: GraphV2) -> None:
+    job_id = gds_arrow_client.get_relationships(sample_graph.name(), ["REL"], relationship_properties=["relX", "relY"])
+    result = gds_arrow_client.stream(sample_graph.name(), job_id)
+
+    assert len(result) == 2
+    assert "sourceNodeId" in result.columns
+    assert "targetNodeId" in result.columns
+    assert "relationshipType" in result.columns
+    assert "REL" in result["relationshipType"].tolist()
+    assert "relX" in result.columns
+    assert "relY" in result.columns
+    assert list(result["relX"].tolist()) == [1.0, 2.0]
+    assert list(result["relY"].tolist()) == [42.0, 43.0]
+
+
+def test_project_from_triplets(arrow_client: AuthenticatedArrowClient, gds_arrow_client: GdsArrowClient) -> None:
+    df = pd.DataFrame(
+        {"sourceNode": np.array([1, 2, 3], dtype=np.int64), "targetNode": np.array([4, 5, 6], dtype=np.int64)}
+    )
+
+    job_id = gds_arrow_client.create_graph_from_triplets("triplets")
+    gds_arrow_client.upload_triplets(job_id, df)
+    gds_arrow_client.triplet_load_done(job_id)
+
+    while gds_arrow_client.job_status(job_id).status != "Done":
+        pass
+
+    listing = CatalogArrowEndpoints(arrow_client).list("triplets")[0]
+    assert listing.node_count == 6
+    assert listing.relationship_count == 3
+    assert listing.graph_name == "triplets"
+
+
+def test_project_from_tables(arrow_client: AuthenticatedArrowClient, gds_arrow_client: GdsArrowClient) -> None:
+    nodes = pd.DataFrame(
+        {
+            "nodeId": np.array([1, 2, 3, 4, 5, 6], dtype=np.int64),
+        }
+    )
+
+    rels = pd.DataFrame(
+        {
+            "sourceNodeId": np.array([1, 2, 3], dtype=np.int64),
+            "targetNodeId": np.array([4, 5, 6], dtype=np.int64),
+        }
+    )
+
+    job_id = gds_arrow_client.create_graph("table")
+    gds_arrow_client.upload_nodes(job_id, nodes)
+    gds_arrow_client.node_load_done(job_id)
+
+    while gds_arrow_client.job_status(job_id).status != "RELATIONSHIP_LOADING":
+        pass
+
+    gds_arrow_client.upload_relationships(job_id, rels)
+    gds_arrow_client.relationship_load_done(job_id)
+
+    while gds_arrow_client.job_status(job_id).status != "Done":
+        pass
+
+    listing = CatalogArrowEndpoints(arrow_client).list("table")[0]
+    assert listing.node_count == 6
+    assert listing.relationship_count == 3
+    assert listing.graph_name == "table"

graphdatascience/tests/integrationV2/conftest.py

@@ -1,8 +1,19 @@
+import logging
 import os
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Generator
 
 import pytest
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.network import Network
+from testcontainers.core.waiting_utils import wait_for_logs
+
+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
+
+LOGGER = logging.getLogger(__name__)
 
 
 def pytest_collection_modifyitems(config: Any, items: Any) -> None:
@@ -31,3 +42,89 @@ def logs_dir(tmp_path_factory: pytest.TempPathFactory) -> Generator[Path, None,
 
 def inside_ci() -> bool:
     return os.environ.get("BUILD_NUMBER") is not None
+
+
+@dataclass
+class GdsSessionConnectionInfo:
+    host: str
+    arrow_port: int
+    bolt_port: int
+
+
+@pytest.fixture(scope="package")
+def password_dir(tmp_path_factory: pytest.TempPathFactory) -> Generator[Path, None, None]:
+    """Create a temporary file and return its path."""
+    tmp_dir = tmp_path_factory.mktemp("passwords")
+    temp_file_path = os.path.join(tmp_dir, "password")
+
+    with open(temp_file_path, "w") as f:
+        f.write("password")
+
+    yield tmp_dir
+
+    # Clean up the file
+    os.unlink(temp_file_path)
+
+
+@pytest.fixture(scope="package")
+def network() -> Generator[Network, None, None]:
+    with Network() as network:
+        yield network
+
+
+def start_session(
+    logs_dir: Path, network: Network, password_dir: Path
+) -> Generator[GdsSessionConnectionInfo, None, None]:
+    if (session_uri := os.environ.get("GDS_SESSION_URI")) is not None:
+        uri_parts = session_uri.split(":")
+        yield GdsSessionConnectionInfo(host=uri_parts[0], arrow_port=8491, bolt_port=int(uri_parts[1]))
+        return
+
+    session_image = os.getenv(
+        "GDS_SESSION_IMAGE", "europe-west1-docker.pkg.dev/gds-aura-artefacts/gds/gds-session:latest"
+    )
+    LOGGER.info(f"Using session image: {session_image}")
+    session_container = (
+        DockerContainer(
+            image=session_image,
+        )
+        .with_env("ALLOW_LIST", "DEFAULT")
+        .with_env("DNS_NAME", "gds-session")
+        .with_env("PAGE_CACHE_SIZE", "100M")
+        .with_env("MODEL_STORAGE_BASE_LOCATION", "/models")
+        .with_exposed_ports(8491)
+        .with_volume_mapping(password_dir, "/passwords")
+    )
+    if not inside_ci():
+        session_container = session_container.with_network(network).with_network_aliases("gds-session")
+    with session_container as session_container:
+        try:
+            wait_for_logs(session_container, "Running GDS tasks: 0", timeout=20)
+            yield GdsSessionConnectionInfo(
+                host=session_container.get_container_host_ip(),
+                arrow_port=session_container.get_exposed_port(8491),
+                bolt_port=-1,  # not used in tests
+            )
+        finally:
+            stdout, stderr = session_container.get_logs()
+
+            if stderr:
+                print(f"Error logs from session container:\n{stderr}")
+
+            if inside_ci():
+                print(f"Session container logs:\n{stdout}")
+
+            out_file = logs_dir / "session_container.log"
+            with open(out_file, "w") as f:
+                f.write(stdout.decode("utf-8"))
+
+
+def create_arrow_client(session_uri: GdsSessionConnectionInfo) -> AuthenticatedArrowClient:
+    """Create an authenticated Arrow client connected to the session container."""
+
+    return AuthenticatedArrowClient.create(
+        arrow_info=ArrowInfo(f"{session_uri.host}:{session_uri.arrow_port}", True, True, ["v1", "v2"]),
+        auth=UsernamePasswordAuthentication("neo4j", "password"),
+        encrypted=False,
+        advertised_listen_address=("gds-session", 8491),
+    )

graphdatascience/tests/integrationV2/procedure_surface/arrow/conftest.py

@@ -9,13 +9,15 @@
 from graphdatascience import QueryRunner
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.session.dbms_connection_info import DbmsConnectionInfo
-from graphdatascience.tests.integrationV2.conftest import inside_ci
-from graphdatascience.tests.integrationV2.procedure_surface.conftest import (
+from graphdatascience.tests.integrationV2.conftest import (
     GdsSessionConnectionInfo,
     create_arrow_client,
+    inside_ci,
+    start_session,
+)
+from graphdatascience.tests.integrationV2.procedure_surface.conftest import (
     create_db_query_runner,
     start_database,
-    start_session,
 )
 
 LOGGER = logging.getLogger(__name__)