From 076ba23cac942f0bfec2a76b8f2ad97725f9527e Mon Sep 17 00:00:00 2001
From: Jordan Matelsky <j6k4m8@gmail.com>
Date: Fri, 24 Jul 2026 14:54:47 -0400
Subject: [PATCH] fix: improve graph file ingestion

---
 .../host_provider/SingleFileHostProvider.py   |  8 +++++-
 .../temporary_graph_host_provider.py          | 25 ++++++-------------
 .../test_filesystem_host_provider.py          |  9 +++++++
 .../test_temporary_graph_host_provider.py     | 11 ++++++++
 4 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/server/src/host_provider/host_provider/SingleFileHostProvider.py b/server/src/host_provider/host_provider/SingleFileHostProvider.py
index 5a4a293..e65dce7 100644
--- a/server/src/host_provider/host_provider/SingleFileHostProvider.py
+++ b/server/src/host_provider/host_provider/SingleFileHostProvider.py
@@ -56,4 +56,10 @@ def get_networkx_graph(self, uri: str) -> nx.Graph:
             nx.Graph: The NetworkX graph.
 
         """
-        return ACCEPTED_EXTENSIONS[uri.split(".")[-1]](uri)
+        extension = next(
+            (candidate for candidate in sorted(ACCEPTED_EXTENSIONS, key=len, reverse=True) if uri.endswith(candidate)),
+            None,
+        )
+        if extension is None:
+            raise ValueError(f"Unsupported graph file: {uri}")
+        return ACCEPTED_EXTENSIONS[extension](uri)
diff --git a/server/src/host_provider/host_provider/temporary_graph_host_provider.py b/server/src/host_provider/host_provider/temporary_graph_host_provider.py
index d1694c8..6e6dd53 100644
--- a/server/src/host_provider/host_provider/temporary_graph_host_provider.py
+++ b/server/src/host_provider/host_provider/temporary_graph_host_provider.py
@@ -211,26 +211,15 @@ def _read_csv_edgelist(self, filepath: str) -> nx.Graph:
             if len(df.columns) < 2:
                 raise ValueError("CSV file must have at least 2 columns for source and target nodes")
 
-            # Create graph from edgelist
-            G = nx.Graph()
-
-            # Get column names
             source_col = df.columns[0]
             target_col = df.columns[1]
-
-            # Add edges
-            for _, row in df.iterrows():
-                source = row[source_col]
-                target = row[target_col]
-
-                # Add edge with any additional attributes
-                edge_attrs = {}
-                for col in df.columns[2:]:
-                    edge_attrs[col] = row[col]
-
-                G.add_edge(source, target, **edge_attrs)
-
-            return G
+            return nx.from_pandas_edgelist(
+                df,
+                source=source_col,
+                target=target_col,
+                edge_attr=list(df.columns[2:]) or None,
+                create_using=nx.Graph,
+            )
 
         except Exception as e:
             # If CSV reading fails, try as plain text edgelist
diff --git a/server/src/host_provider/host_provider/test_filesystem_host_provider.py b/server/src/host_provider/host_provider/test_filesystem_host_provider.py
index 9759351..c39c24a 100644
--- a/server/src/host_provider/host_provider/test_filesystem_host_provider.py
+++ b/server/src/host_provider/host_provider/test_filesystem_host_provider.py
@@ -27,3 +27,12 @@ def test_can_count_motifs():
         r = FilesystemGraphHostProvider()
         # Query the provider for the motif count.
         assert r.get_motif_count(uri, "A->B") == 12
+
+
+def test_can_read_compressed_graphml():
+    g = nx.path_graph(4)
+    with tempfile.NamedTemporaryFile(suffix=".graphml.gz") as f:
+        nx.write_graphml(g, f.name)
+        provider = FilesystemGraphHostProvider()
+
+        assert nx.is_isomorphic(g, provider.get_networkx_graph("file://" + f.name))
diff --git a/server/src/host_provider/host_provider/test_temporary_graph_host_provider.py b/server/src/host_provider/host_provider/test_temporary_graph_host_provider.py
index 1ebb9c7..fcb6c38 100644
--- a/server/src/host_provider/host_provider/test_temporary_graph_host_provider.py
+++ b/server/src/host_provider/host_provider/test_temporary_graph_host_provider.py
@@ -37,3 +37,14 @@ def test_display_name_survives_provider_restart(tmp_path):
     restarted = TemporaryGraphHostProvider(str(tmp_path))
 
     assert restarted.get_file_info(temp_id)["display_name"] == "My graph"
+
+
+def test_csv_attributes_are_preserved(tmp_path):
+    provider = TemporaryGraphHostProvider(str(tmp_path))
+    path = tmp_path / "attributes.csv"
+    path.write_text("source,target,weight,label\na,b,2,edge\n")
+
+    graph = provider._read_csv_edgelist(str(path))
+
+    assert graph.edges["a", "b"]["weight"] == 2
+    assert graph.edges["a", "b"]["label"] == "edge"