Allow Driver as argument to from_neo4j

Author: adamnsch

changelog.md

@@ -8,6 +8,8 @@
 
 ## New features
 
+* Allow passing a `neo4j.Driver` instance as input to `from_neo4j`, in which case the driver will be used internally to fetch the graph data using a simple query
+
 
 ## Bug fixes
 

docs/source/integration.rst

@@ -184,13 +184,16 @@ Once you have installed the additional dependency, you can use the :doc:`from_ne
 to import query results from Neo4j.
 
 The ``from_neo4j`` method takes one mandatory positional parameter:
-
-* A ``result`` representing the query result either in form of `neo4j.graph.Graph` or `neo4j.Result`.
+A ``data`` argument representing either a query result in the shape of a ``neo4j.graph.Graph`` or ``neo4j.Result``, or a
+``neo4j.Driver`` in which case a simple default query will be executed internally to retrieve the graph data.
 
 We can also provide an optional ``size_property`` parameter, which should refer to a node property,
 and will be used to determine the sizes of the nodes in the visualization.
 
-The ``node_caption`` and ``relationship_caption`` parameters are also optional, and indicate the node and relationship properties to use for the captions of each element in the visualization.
+The ``node_caption`` and ``relationship_caption`` parameters are also optional, and indicate the node and relationship
+properties to use for the captions of each element in the visualization.
+By default, the captions will be set to the node labels relationship types, but you can specify any property that
+exists on these entities.
 
 The last optional property, ``node_radius_min_max``, can be used (and is used by default) to scale the node sizes for
 the visualization.

python-wrapper/src/neo4j_viz/neo4j.py

@@ -3,7 +3,7 @@
 from typing import Optional, Union
 
 import neo4j.graph
-from neo4j import Result
+from neo4j import Driver, Result, RoutingControl
 from pydantic import BaseModel, ValidationError
 
 from neo4j_viz.node import Node
@@ -20,14 +20,15 @@ def _parse_validation_error(e: ValidationError, entity_type: type[BaseModel]) ->
 
 
 def from_neo4j(
-    result: Union[neo4j.graph.Graph, Result],
+    data: Union[neo4j.graph.Graph, Result, Driver],
     size_property: Optional[str] = None,
     node_caption: Optional[str] = "labels",
     relationship_caption: Optional[str] = "type",
     node_radius_min_max: Optional[tuple[float, float]] = (3, 60),
+    row_limit: int = 10_000,
 ) -> VisualizationGraph:
     """
-    Create a VisualizationGraph from a Neo4j Graph or Neo4j Result object.
+    Create a VisualizationGraph from a Neo4j `Graph`, Neo4j `Result` or Neo4j `Driver`.
 
     All node and relationship properties will be included in the visualization graph.
     If the properties are named as the fields of the `Node` or `Relationship` classes, they will be included as
@@ -36,8 +37,9 @@ def from_neo4j(
 
     Parameters
     ----------
-    result : Union[neo4j.graph.Graph, Result]
-        Query result either in shape of a Graph or result.
+    data : Union[neo4j.graph.Graph, neo4j.Result, neo4j.Driver]
+        Either a query result in the shape of a `neo4j.graph.Graph` or `neo4j.Result`, or a `neo4j.Driver` in
+        which case a simple default query will be executed internally to retrieve the graph data.
     size_property : str, optional
         Property to use for node size, by default None.
     node_caption : str, optional
@@ -47,14 +49,23 @@ def from_neo4j(
     node_radius_min_max : tuple[float, float], optional
         Minimum and maximum node radius, by default (3, 60).
         To avoid tiny or huge nodes in the visualization, the node sizes are scaled to fit in the given range.
+    row_limit : int, optional
+        Maximum number of rows to return from the query, by default 10_000.
+        This is only used if a `neo4j.Driver` is passed as `result` argument, otherwise the limit is ignored.
     """
 
-    if isinstance(result, Result):
-        graph = result.graph()
-    elif isinstance(result, neo4j.graph.Graph):
-        graph = result
+    if isinstance(data, Result):
+        graph = data.graph()
+    elif isinstance(data, neo4j.graph.Graph):
+        graph = data
+    elif isinstance(data, Driver):
+        graph = data.execute_query(
+            f"MATCH (n)-[r]->(m) RETURN n,r,m LIMIT {row_limit}",
+            routing_=RoutingControl.READ,
+            result_transformer_=Result.graph,
+        )
     else:
-        raise ValueError(f"Invalid input type `{type(result)}`. Expected `neo4j.Graph` or `neo4j.Result`")
+        raise ValueError(f"Invalid input type `{type(data)}`. Expected `neo4j.Graph`, `neo4j.Result` or `neo4j.Driver`")
 
     all_node_field_aliases = Node.all_validation_aliases()
     all_rel_field_aliases = Relationship.all_validation_aliases()

python-wrapper/tests/test_neo4j.py

@@ -2,7 +2,7 @@
 
 import neo4j
 import pytest
-from neo4j import Session
+from neo4j import Driver, Session
 
 from neo4j_viz.neo4j import from_neo4j
 from neo4j_viz.node import Node
@@ -201,3 +201,51 @@ def test_from_neo4j_rel_error(neo4j_session: Session) -> None:
         match="Error for relationship property 'caption_align' with provided input 'banana'. Reason: Input should be 'top', 'center' or 'bottom'",
     ):
         from_neo4j(graph)
+
+
+@pytest.mark.requires_neo4j_and_gds
+def test_from_neo4j_graph_driver(neo4j_session: Session, neo4j_driver: Driver) -> None:
+    graph = neo4j_session.run("MATCH (a:_CI_A|_CI_B)-[r]->(b) RETURN a, b, r ORDER BY a").graph()
+
+    # Note that this tests requires an empty Neo4j database, as it just fetches everything
+    VG = from_neo4j(neo4j_driver)
+
+    sorted_nodes: list[neo4j.graph.Node] = sorted(graph.nodes, key=lambda x: dict(x.items())["name"])
+    node_ids: list[str] = [node.element_id for node in sorted_nodes]
+
+    expected_nodes = [
+        Node(
+            id=node_ids[0],
+            caption="_CI_A",
+            properties=dict(
+                labels=["_CI_A"],
+                name="Alice",
+                height=20,
+                id=42,
+                _id=1337,
+                caption="hello",
+            ),
+        ),
+        Node(
+            id=node_ids[1],
+            caption="_CI_A:_CI_B",
+            size=11,
+            properties=dict(
+                labels=["_CI_A", "_CI_B"],
+                name="Bob",
+                height=10,
+                id=84,
+                __labels=[1, 2],
+            ),
+        ),
+    ]
+
+    assert len(VG.nodes) == 2
+    assert sorted(VG.nodes, key=lambda x: x.properties["name"]) == expected_nodes
+
+    assert len(VG.relationships) == 2
+    vg_rels = sorted([(e.source, e.target, e.caption) for e in VG.relationships], key=lambda x: x[2] if x[2] else "foo")
+    assert vg_rels == [
+        (node_ids[0], node_ids[1], "KNOWS"),
+        (node_ids[1], node_ids[0], "RELATED"),
+    ]