diff --git a/.github/workflows/gds-integration-tests.yml b/.github/workflows/gds-integration-tests.yml index 53a94b8..69a40f3 100644 --- a/.github/workflows/gds-integration-tests.yml +++ b/.github/workflows/gds-integration-tests.yml @@ -20,6 +20,7 @@ jobs: fail-fast: false matrix: os: [ubuntu-latest] + gds-version: ["1.22", "2.0.0a1"] defaults: run: working-directory: python-wrapper @@ -33,11 +34,12 @@ jobs: with: python-version: "3.11" enable-cache: true - - run: uv sync --group dev --extra pandas --extra neo4j --extra gds + - name: Install just + uses: extractions/setup-just@v2 - name: Run tests env: AURA_API_CLIENT_ID: 4V1HYCYEeoU4dSxThKnBeLvE2U4hSphx AURA_API_CLIENT_SECRET: ${{ secrets.AURA_API_CLIENT_SECRET }} AURA_API_PROJECT_ID: 3f8df5e7-4800-4d4f-ad1d-2d044dfd587c - run: uv run pytest tests/ --include-neo4j-and-gds + run: just py-ci-test-gds "${{ matrix.gds-version }}" diff --git a/justfile b/justfile index 21446a7..7a8a995 100644 --- a/justfile +++ b/justfile @@ -25,6 +25,16 @@ py-test: cd python-wrapper && uv sync --all-extras --group dev cd python-wrapper && uv run --group dev pytest +# install a specific GDS client version and run the GDS integration tests (used by CI) +# example: just py-ci-test-gds 2.0.0a1 +py-ci-test-gds gds_version: + #!/usr/bin/env bash + set -e + cd {{py_dir}} + uv sync --group dev --extra pandas --extra neo4j --extra gds + uv pip install "graphdatascience=={{gds_version}}" + uv run pytest tests/ --include-neo4j-and-gds + py-test-gds: #!/usr/bin/env bash set -e diff --git a/python-wrapper/pyproject.toml b/python-wrapper/pyproject.toml index 986119d..330bbd7 100644 --- a/python-wrapper/pyproject.toml +++ b/python-wrapper/pyproject.toml @@ -43,7 +43,7 @@ requires-python = ">=3.10" [project.optional-dependencies] pandas = ["pandas>=2, <3", "pandas-stubs>=2, <3"] -gds = ["graphdatascience>=1.22, <2"] +gds = ["graphdatascience>=1.22, <3"] neo4j = ["neo4j"] snowflake = ["snowflake-snowpark-python>=1, <2"] diff --git a/python-wrapper/src/neo4j_viz/_gds_compat.py b/python-wrapper/src/neo4j_viz/_gds_compat.py new file mode 100644 index 0000000..d09ccb8 --- /dev/null +++ b/python-wrapper/src/neo4j_viz/_gds_compat.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import importlib +import re +from typing import Any + +from graphdatascience.version import __version__ as _gds_version + + +def _parse_major(version: str) -> int: + match = re.match(r"\s*(\d+)", version) + return int(match.group(1)) if match else 0 + + +IS_GDS_2: bool = _parse_major(_gds_version) >= 2 + +if IS_GDS_2: + GdsGraph: Any = importlib.import_module("graphdatascience.graph").Graph + _GRAPH_TYPES: tuple[type, ...] = (GdsGraph,) +else: + GdsGraph = importlib.import_module("graphdatascience.graph.v2").GraphV2 + _GRAPH_TYPES = (GdsGraph, importlib.import_module("graphdatascience").Graph) + + +def _check_graph_type(G: Any) -> None: + """Raise ``TypeError`` unless ``G`` is a graph object accepted by the installed client.""" + if not isinstance(G, _GRAPH_TYPES): + accepted = " or ".join(t.__name__ for t in _GRAPH_TYPES) + raise TypeError(f"`G` must be a GDS graph object ({accepted}), but got {type(G).__name__}") + + +def _catalog(gds: Any) -> Any: + """Return the graph catalog endpoints for either client version.""" + return gds.graph if IS_GDS_2 else gds.v2.graph + + +def _degree_centrality(gds: Any) -> Any: + """Return the degree centrality endpoints for either client version.""" + return gds.degree_centrality if IS_GDS_2 else gds.v2.degree_centrality diff --git a/python-wrapper/src/neo4j_viz/gds.py b/python-wrapper/src/neo4j_viz/gds.py index 78e37ba..67a1a42 100644 --- a/python-wrapper/src/neo4j_viz/gds.py +++ b/python-wrapper/src/neo4j_viz/gds.py @@ -2,29 +2,29 @@ import warnings from itertools import chain -from typing import Collection, Optional +from typing import Any, Collection, Optional from uuid import uuid4 import pandas as pd -from graphdatascience import Graph, GraphDataScience -from graphdatascience.graph.v2 import GraphV2 +from graphdatascience import GraphDataScience from graphdatascience.session import AuraGraphDataScience from neo4j_viz.colors import NEO4J_COLORS_DISCRETE, ColorSpace +from ._gds_compat import IS_GDS_2, GdsGraph, _catalog, _check_graph_type, _degree_centrality from .pandas import _from_dfs from .visualization_graph import VisualizationGraph def _fetch_node_dfs( gds: GraphDataScience | AuraGraphDataScience, - G: GraphV2, + G: Any, node_properties_by_label: dict[str, list[str]], node_labels: Collection[str], additional_db_node_properties: list[str], ) -> dict[str, pd.DataFrame]: return { - lbl: gds.v2.graph.node_properties.stream( + lbl: _catalog(gds).node_properties.stream( G, node_properties=node_properties_by_label[lbl], node_labels=[lbl], @@ -34,14 +34,14 @@ def _fetch_node_dfs( } -def _fetch_rel_dfs(gds: GraphDataScience | AuraGraphDataScience, G: GraphV2) -> list[pd.DataFrame]: +def _fetch_rel_dfs(gds: GraphDataScience | AuraGraphDataScience, G: Any) -> list[pd.DataFrame]: rel_props = G.relationship_properties() rel_dfs: list[pd.DataFrame] = [] # Have to call per stream per relationship type as there was a bug in GDS < 2.21 for rel_type, props in rel_props.items(): - rel_df = gds.v2.graph.relationships.stream( + rel_df = _catalog(gds).relationships.stream( G, relationship_types=[rel_type], relationship_properties=list(props) ) @@ -62,7 +62,7 @@ def _fetch_rel_dfs(gds: GraphDataScience | AuraGraphDataScience, G: GraphV2) -> def from_gds( gds: GraphDataScience | AuraGraphDataScience, - G: Graph | GraphV2, + G: Any, node_properties: Optional[list[str]] = None, db_node_properties: Optional[list[str]] = None, max_node_count: int = 10_000, @@ -97,8 +97,9 @@ def from_gds( """ if db_node_properties is None: db_node_properties = [] - if isinstance(G, Graph): - G_v2 = gds.v2.graph.get(G.name()) + _check_graph_type(G) + if not IS_GDS_2 and not isinstance(G, GdsGraph): + G_v2 = _catalog(gds).get(G.name()) else: G_v2 = G @@ -127,7 +128,7 @@ def from_gds( ) sampling_ratio = float(max_node_count) / node_count sample_name = f"neo4j-viz_sample_{uuid4()}" - G_fetched, _ = gds.v2.graph.sample.rwr( + G_fetched, _ = _catalog(gds).sample.rwr( G_v2, sample_name, sampling_ratio=sampling_ratio, node_label_stratification=True ) else: @@ -139,7 +140,7 @@ def from_gds( # as a temporary property to ensure that we have at least one property for each label to fetch if sum([len(props) == 0 for props in node_properties_by_label.values()]) > 0: property_name = f"neo4j-viz_property_{uuid4()}" - gds.v2.degree_centrality.mutate(G_fetched, mutate_property=property_name) + _degree_centrality(gds).mutate(G_fetched, mutate_property=property_name) for props in node_properties_by_label.values(): props.append(property_name) @@ -155,7 +156,7 @@ def from_gds( if G_fetched.name() != G.name(): G_fetched.drop() elif property_name is not None: - gds.v2.graph.node_properties.drop(G_fetched, node_properties=[property_name]) + _catalog(gds).node_properties.drop(G_fetched, node_properties=[property_name]) for df in node_dfs.values(): if property_name is not None and property_name in df.columns: diff --git a/python-wrapper/tests/neo4j_and_gds/test_gds.py b/python-wrapper/tests/neo4j_and_gds/test_gds.py index 46d9994..dfbdeab 100644 --- a/python-wrapper/tests/neo4j_and_gds/test_gds.py +++ b/python-wrapper/tests/neo4j_and_gds/test_gds.py @@ -1,14 +1,13 @@ import re -from contextlib import AbstractContextManager -from typing import Generator +from typing import Any, Generator import pandas as pd import pytest from graphdatascience import GraphDataScience -from graphdatascience.graph.v2 import GraphV2 from graphdatascience.session import AuraGraphDataScience from neo4j_viz import Node +from neo4j_viz._gds_compat import _catalog from neo4j_viz.gds import from_gds @@ -25,11 +24,11 @@ def db_setup(gds: GraphDataScience | AuraGraphDataScience) -> Generator[None, No gds.run_cypher("MATCH (n:_CI_A|_CI_B) DETACH DELETE n") -def project_graph(gds: GraphDataScience | AuraGraphDataScience) -> AbstractContextManager[GraphV2]: +def project_graph(gds: GraphDataScience | AuraGraphDataScience) -> Any: if isinstance(gds, GraphDataScience): - return gds.v2.graph.project("g2", "*", "*") + return _catalog(gds).project("g2", "*", "*") elif isinstance(gds, AuraGraphDataScience): - return gds.v2.graph.project("g2", "MATCH (n)–->(m) RETURN gds.graph.project.remote(n, m)") + return _catalog(gds).project("g2", "MATCH (n)–->(m) RETURN gds.graph.project.remote(n, m)") raise Exception(f"Unsupported GDS type {type(gds)}") @@ -64,7 +63,7 @@ def test_from_gds_integration_all_properties(gds: GraphDataScience | AuraGraphDa } ) - with gds.v2.graph.construct("flo", nodes, rels) as G: + with _catalog(gds).construct("flo", nodes, rels) as G: VG = from_gds(gds, G) assert len(VG.nodes) == 3 @@ -113,7 +112,7 @@ def test_from_gds_integration_all_properties(gds: GraphDataScience | AuraGraphDa @pytest.mark.requires_neo4j_and_gds def test_from_gds_sample(gds: GraphDataScience | AuraGraphDataScience) -> None: - with gds.v2.graph.generate("hello", node_count=11_000, average_degree=1) as G: + with _catalog(gds).generate("hello", node_count=11_000, average_degree=1) as G: with pytest.warns( UserWarning, match=re.escape( @@ -164,7 +163,7 @@ def test_from_gds_hetero(gds: GraphDataScience | AuraGraphDataScience) -> None: } ) - with gds.v2.graph.construct("flo", [A_nodes, B_nodes], [X_rels, Y_rels]) as G: + with _catalog(gds).construct("flo", [A_nodes, B_nodes], [X_rels, Y_rels]) as G: VG = from_gds( gds, G,