diff --git a/grand/backends/_sqlbackend.py b/grand/backends/_sqlbackend.py index af20e23..67a65bf 100644 --- a/grand/backends/_sqlbackend.py +++ b/grand/backends/_sqlbackend.py @@ -776,20 +776,46 @@ 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( @@ -797,26 +823,33 @@ def ingest_from_edgelist_dataframe( 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), diff --git a/grand/backends/test_sql_transactions.py b/grand/backends/test_sql_transactions.py index 1ffad3c..7406829 100644 --- a/grand/backends/test_sql_transactions.py +++ b/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()