Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 36 additions & 11 deletions grand/backends/_sqlbackend.py
Original file line number Diff line number Diff line change
Expand Up @@ -373,18 +373,43 @@ def add_edge(self, u: Hashable, v: Hashable, metadata: dict):
return pk

def add_edges_from(self, ebunch_to_add, **attr):
edges = [
{
self._primary_key: f"__{u}__{v}",
self._edge_source_key: u,
self._edge_target_key: v,
"_metadata": {**attr, **metadata},
}
for u, v, metadata in ebunch_to_add
]
edges = []
updates = []
endpoints = set()
for edge in ebunch_to_add:
if len(edge) == 2:
u, v = edge
metadata = {}
else:
u, v, metadata = edge
metadata = {**attr, **metadata}
endpoints.update((u, v))
if self.has_edge(u, v):
existing_metadata = self.get_edge_by_id(u, v)
existing_metadata.update(metadata)
updates.append((u, v, existing_metadata))
else:
edges.append(
{
self._primary_key: f"__{u}__{v}",
self._edge_source_key: u,
self._edge_target_key: v,
"_metadata": metadata,
}
)

with self._mutation():
self._connection.execute(self._edge_table.insert(), edges)
with self.transaction():
for node in endpoints:
self._insert_empty_node_if_missing(node)
if edges:
self._connection.execute(self._edge_table.insert(), edges)
for u, v, metadata in updates:
self._connection.execute(
self._edge_table.update().where(
self._edge_table.c[self._primary_key] == f"__{u}__{v}"
),
parameters={"_metadata": metadata},
)

def all_edges_as_iterable(self, include_metadata: bool = False) -> Generator:
"""
Expand Down
44 changes: 44 additions & 0 deletions grand/backends/test_sql_transactions.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,50 @@ def test_transaction_context_rolls_back_grouped_mutations(tmp_path):
backend.close()


def test_add_edges_from_creates_endpoints_and_accepts_two_tuples(tmp_path):
backend = SQLBackend(db_url=f"sqlite:///{tmp_path / 'graph.db'}", directed=True)

backend.add_edges_from([("A", "B"), ("B", "C", {"weight": 2})])

assert set(backend.all_nodes_as_iterable()) == {"A", "B", "C"}
assert backend.has_edge("A", "B")
assert backend.get_edge_by_id("B", "C") == {"weight": 2}
backend.close()


def test_add_edges_from_merges_existing_edge_metadata(tmp_path):
backend = SQLBackend(db_url=f"sqlite:///{tmp_path / 'graph.db'}", directed=True)
backend.add_edge("A", "B", {"weight": 1, "kind": "existing"})

backend.add_edges_from([("A", "B", {"weight": 2})], batch=True)

assert backend.get_edge_by_id("A", "B") == {
"weight": 2,
"kind": "existing",
"batch": True,
}
backend.close()


def test_add_edges_from_rolls_back_entire_batch_on_failure(tmp_path):
backend = SQLBackend(db_url=f"sqlite:///{tmp_path / 'graph.db'}", directed=True)
original_execute = backend._connection.execute
def fail_edge_insert(statement, *args, **kwargs):
if getattr(statement, "table", None) is backend._edge_table:
raise RuntimeError("edge batch failed")
return original_execute(statement, *args, **kwargs)

backend._connection.execute = fail_edge_insert

with pytest.raises(RuntimeError, match="edge batch failed"):
backend.add_edges_from([("A", "B"), ("B", "C")])

backend._connection.execute = original_execute
assert backend.get_node_count() == 0
assert backend.get_edge_count() == 0
backend.close()


def test_add_edge_rolls_back_created_nodes_when_edge_insert_fails(tmp_path):
backend = SQLBackend(db_url=f"sqlite:///{tmp_path / 'graph.db'}", directed=True)
original_execute = backend._connection.execute
Expand Down
Loading