fix: preserve data during SQL ingestion

grand/backends/_sqlbackend.py

@@ -776,47 +776,80 @@ def ingest_from_edgelist_dataframe(
 
         edge_rows = [
             {
+                self._primary_key: f"__{source}__{target}",
                 self._edge_source_key: source,
                 self._edge_target_key: target,
-                self._primary_key: f"__{source}__{target}",
                 "_metadata": metadata,
             }
             for source, target, metadata in zip(sources, targets, edge_metadata)
         ]
+        edge_ids = [row[self._primary_key] for row in edge_rows]
 
-        with self._mutation():
-            if edge_rows:
-                self._connection.execute(self._edge_table.insert(), edge_rows)
-
+        with self.transaction():
+            try:
+                with self._connection.begin_nested():
+                    if edge_rows:
+                        self._connection.execute(self._edge_table.insert(), edge_rows)
+            except sqlalchemy.exc.IntegrityError:
+                existing = {
+                    row[self._primary_key]: row["_metadata"]
+                    for row in self._connection.execute(
+                        select(
+                            self._edge_table.c[self._primary_key],
+                            self._edge_table.c["_metadata"],
+                        ).where(self._edge_table.c[self._primary_key].in_(edge_ids))
+                    ).mappings()
+                }
+                new_edges = [
+                    row for row in edge_rows if row[self._primary_key] not in existing
+                ]
+                if new_edges:
+                    self._connection.execute(self._edge_table.insert(), new_edges)
+                for row in edge_rows:
+                    edge_id = row[self._primary_key]
+                    if edge_id in existing:
+                        metadata = {**existing[edge_id], **row["_metadata"]}
+                        self._connection.execute(
+                            self._edge_table.update().where(
+                                self._edge_table.c[self._primary_key] == edge_id
+                            ),
+                            parameters={"_metadata": metadata},
+                        )
             edge_toc = time.time() - edge_tic
-
             node_tic = time.time()
             nodes = pd.unique(
                 pd.concat(
                     [edgelist[source_column], edgelist[target_column]],
                     ignore_index=True,
                 )
             )
-
             node_rows = [
-                {
-                    self._primary_key: str(node),
-                    "_metadata": {},
-                }
-                for node in nodes
+                {self._primary_key: str(node), "_metadata": {}} for node in nodes
             ]
-
             if node_rows:
-                node_insert = self._node_table.insert()
+                insert = self._node_table.insert()
                 if self._engine.dialect.name == "sqlite":
-                    node_insert = node_insert.prefix_with("OR IGNORE")
-                    self._connection.execute(node_insert, node_rows)
+                    insert = insert.prefix_with("OR IGNORE")
                 elif self._engine.dialect.name in {"mysql", "mariadb"}:
-                    node_insert = node_insert.prefix_with("IGNORE")
-                    self._connection.execute(node_insert, node_rows)
+                    insert = insert.prefix_with("IGNORE")
                 else:
-                    for node in nodes:
-                        self._insert_empty_node_if_missing(node)
+                    existing_nodes = set(
+                        self._connection.execute(
+                            select(self._node_table.c[self._primary_key]).where(
+                                self._node_table.c[self._primary_key].in_(
+                                    [row[self._primary_key] for row in node_rows]
+                                )
+                            )
+                        ).scalars()
+                    )
+                    node_rows = [
+                        row
+                        for row in node_rows
+                        if row[self._primary_key] not in existing_nodes
+                    ]
+                if node_rows:
+                    self._connection.execute(insert, node_rows)
+
 
         return {
             "node_count": len(nodes),

grand/backends/test_sql_transactions.py

@@ -1,6 +1,7 @@
 import pytest
 from concurrent.futures import ThreadPoolExecutor
 from unittest.mock import Mock
+import pandas as pd
 
 sqlalchemy = pytest.importorskip("sqlalchemy")
 
@@ -174,3 +175,76 @@ def test_concurrent_mutations_do_not_share_connection_simultaneously(tmp_path):
 
     assert backend.get_node_count() == 20
     backend.close()
+
+
+def test_ingest_preserves_existing_nodes_metadata_and_edges(tmp_path):
+    backend = SQLBackend(
+        db_url=f"sqlite:///{tmp_path / 'graph.db'}",
+        directed=True,
+    )
+    backend.add_node("existing", {"kind": "preserved"})
+    backend.add_edge("existing", "A", {"weight": 1, "label": "old"})
+    edgelist = pd.DataFrame(
+        {
+            "source": ["existing", "A"],
+            "target": ["A", "B"],
+            "weight": [2, 3],
+        }
+    )
+
+    backend.ingest_from_edgelist_dataframe(edgelist, "source", "target")
+
+    assert backend.get_node_by_id("existing") == {"kind": "preserved"}
+    assert backend.get_edge_by_id("existing", "A") == {
+        "weight": 2,
+        "label": "old",
+    }
+    assert backend.get_edge_by_id("A", "B") == {"weight": 3}
+    assert set(backend.all_nodes_as_iterable()) == {"existing", "A", "B"}
+    backend.close()
+
+
+def test_ingest_preserves_primary_key_schema(tmp_path):
+    backend = SQLBackend(db_url=f"sqlite:///{tmp_path / 'graph.db'}")
+
+    backend.ingest_from_edgelist_dataframe(
+        pd.DataFrame({"source": ["A"], "target": ["B"]}),
+        "source",
+        "target",
+    )
+
+    assert backend._node_table.primary_key.columns.keys() == ["ID"]
+    with pytest.raises(sqlalchemy.exc.IntegrityError):
+        backend._connection.execute(
+            backend._node_table.insert(),
+            [
+                {"ID": "duplicate", "_metadata": {}},
+                {"ID": "duplicate", "_metadata": {}},
+            ],
+        )
+    backend._connection.rollback()
+    backend.close()
+
+
+def test_ingest_rolls_back_nodes_and_edges_on_failure(tmp_path):
+    backend = SQLBackend(db_url=f"sqlite:///{tmp_path / 'graph.db'}")
+    original_execute = backend._connection.execute
+
+    def fail_edge_insert(statement, *args, **kwargs):
+        if getattr(statement, "table", None) is backend._edge_table:
+            raise RuntimeError("ingest failed")
+        return original_execute(statement, *args, **kwargs)
+
+    backend._connection.execute = fail_edge_insert
+
+    with pytest.raises(RuntimeError, match="ingest failed"):
+        backend.ingest_from_edgelist_dataframe(
+            pd.DataFrame({"source": ["A"], "target": ["B"]}),
+            "source",
+            "target",
+        )
+
+    backend._connection.execute = original_execute
+    assert backend.get_node_count() == 0
+    assert backend.get_edge_count() == 0
+    backend.close()