Allow overriding caption using dedicated column

Author: s1ck

python-wrapper/src/neo4j_viz/snowflake.py

@@ -318,8 +318,8 @@ def from_snowflake(
 
     By default:
 
-    * The caption of the nodes will be set to the table name.
-    * The caption of the relationships will be set to the table name.
+    * The caption of the nodes will be read from a 'CAPTION' column or set to the table name if the colum does not exist.
+    * The caption of the relationships will be read from a 'CAPTION' column or set to the table name if the column does not exist.
     * The color of the nodes will be set based on the caption, unless there are more than 12 node tables used.
 
     Otherwise, columns will be included as properties on the nodes and relationships.
@@ -334,9 +334,9 @@ def from_snowflake(
     node_dfs, rel_dfs, rel_table_names = _map_tables(session, project_model)
 
     for i, node_df in enumerate(node_dfs):
-        node_df["table"] = project_model.nodeTables[i].split(".")[-1]
+        node_df["table"] = node_df.get("CAPTION", project_model.nodeTables[i].split(".")[-1])
     for i, rel_df in enumerate(rel_dfs):
-        rel_df["table"] = rel_table_names[i].split(".")[-1]
+        rel_df["table"] = rel_df.get("CAPTION", rel_table_names[i].split(".")[-1])
 
     VG = from_dfs(node_dfs, rel_dfs)
 

python-wrapper/tests/test_snowflake.py

@@ -1,7 +1,8 @@
 import pytest
 from snowflake.snowpark import Session
-from snowflake.snowpark.types import LongType, StructField, StructType
+from snowflake.snowpark.types import LongType, StringType, StructField, StructType
 
+from neo4j_viz import Relationship
 from neo4j_viz.node import Node
 from neo4j_viz.snowflake import from_snowflake
 
@@ -60,13 +61,82 @@ def test_from_snowflake(session_with_minimal_graph: Session) -> None:
     )
 
     assert VG.nodes == [
-        Node(id=0, caption="NODES", color="#ffdf81", properties={"SNOWFLAKEID": 6, "table": "NODES"}),
-        Node(id=1, caption="NODES", color="#ffdf81", properties={"SNOWFLAKEID": 7, "table": "NODES"}),
+        Node(id=0, caption="NODES", color="#ffdf81", properties={"SNOWFLAKEID": 6}),
+        Node(id=1, caption="NODES", color="#ffdf81", properties={"SNOWFLAKEID": 7}),
     ]
 
     assert len(VG.relationships) == 1
 
     assert VG.relationships[0].source == 0
     assert VG.relationships[0].target == 1
     assert VG.relationships[0].caption == "RELS"
-    assert VG.relationships[0].properties == {"table": "RELS"}
+    assert VG.relationships[0].properties == {}
+
+
+@pytest.fixture
+def session_with_custom_caption_graph(session: Session) -> Session:
+    """
+    Create a minimal graph with two nodes and one relationship.
+    """
+    node_df = session.create_dataframe(
+        data=[
+            [6, "Alice"],
+            [7, "Bob"],
+            [8, "Charlie"],
+        ],
+        schema=StructType(
+            [
+                StructField("NODEID", LongType()),
+                StructField("CAPTION", StringType()),
+            ]
+        ),
+    )
+    node_df.write.save_as_table("NODES")
+
+    rel_df = session.create_dataframe(
+        data=[
+            [6, 7, "KNOWS"],
+            [6, 8, "LIKES"],
+        ],
+        schema=StructType(
+            [
+                StructField("SOURCENODEID", LongType()),
+                StructField("TARGETNODEID", LongType()),
+                StructField("CAPTION", StringType()),
+            ]
+        ),
+    )
+    rel_df.write.save_as_table("RELS")
+
+    return session
+
+
+def test_from_snowflake_with_caption(session_with_custom_caption_graph: Session) -> None:
+    VG = from_snowflake(
+        session_with_custom_caption_graph,
+        {
+            "nodeTables": ["NODES"],
+            "relationshipTables": {
+                "RELS": {
+                    "sourceTable": "NODES",
+                    "targetTable": "NODES",
+                },
+            },
+        },
+    )
+
+    assert VG.nodes == [
+        Node(id=0, caption="Alice", color="#ffdf81", properties={"CAPTION": "Alice", "SNOWFLAKEID": 6}),
+        Node(id=1, caption="Bob", color="#c990c0", properties={"CAPTION": "Bob", "SNOWFLAKEID": 7}),
+        Node(id=2, caption="Charlie", color="#f79767", properties={"CAPTION": "Charlie", "SNOWFLAKEID": 8}),
+    ]
+
+    assert len(VG.relationships) == 2
+    # ignore uuid in comparison
+    VG.relationships[0].id = 0
+    VG.relationships[1].id = 1
+
+    assert VG.relationships == [
+        Relationship(id=0, source=0, target=1, caption="KNOWS", properties={"CAPTION": "KNOWS"}),
+        Relationship(id=1, source=0, target=2, caption="LIKES", properties={"CAPTION": "LIKES"}),
+    ]