tests/test_graphdb_push_indexes.py
Rama v8 · réplica de lectura del repositorio original.
"""#3804: graph DB pushes must index (label, id) before the MERGE upsert loop.
`push_to_falkordb` / `push_to_neo4j` upsert every node with
MERGE (n:<Label> {id: $id}) SET n += $props
which matches on ``(label, id)``. With no index on ``(n.id)`` the engine has to
scan every node already carrying that label to decide whether the node is new,
so each upsert costs O(nodes with that label) and a full push degrades to
O(n^2). Reported on a 167,428-node / 209,215-edge graph, where throughput
collapsed from ~10,000 node upserts/60s to ~100 upserts/300s partway through.
These tests drive the two push functions with fake SDK clients and assert the
Cypher that reaches the server, so they run without a database.
"""
from __future__ import annotations
import sys
import types
import networkx as nx
import pytest
def _graph() -> nx.Graph:
"""Two labels, so the per-label index creation is observable."""
G = nx.Graph()
G.add_node("parse_config", label="parse_config()", file_type="code", source_file="a.py")
G.add_node("load_settings", label="load_settings()", file_type="code", source_file="a.py")
G.add_node("sftp_service", label="SFTP service", file_type="document", source_file="b.md")
G.add_node("payroll_db", label="Payroll DB", file_type="database", source_file="b.md")
G.add_edge("parse_config", "sftp_service", relation="reads")
G.add_edge("load_settings", "payroll_db", relation="reads")
return G
class _RecordingGraph:
"""Stand-in for the FalkorDB graph handle; records Cypher in call order."""
def __init__(self, fail_on_index: bool = False) -> None:
self.calls: list[str] = []
self._fail_on_index = fail_on_index
def query(self, cypher: str, params: dict | None = None) -> None:
self.calls.append(cypher)
if self._fail_on_index and cypher.startswith("CREATE INDEX"):
# What a re-run against an already-indexed graph gets back:
# FalkorDB's CREATE INDEX has no IF NOT EXISTS.
raise RuntimeError("index already exists")
class _RecordingSession:
def __init__(self, graph: _RecordingGraph) -> None:
self._graph = graph
self.closed = False
def run(self, cypher: str, **params) -> None:
self._graph.calls.append(cypher)
def __enter__(self) -> "_RecordingSession":
return self
def __exit__(self, *exc) -> bool:
self.closed = True
return False
@pytest.fixture()
def fake_falkordb(monkeypatch):
"""Install a fake `falkordb` module and return its recording graph handle."""
graph = _RecordingGraph()
client = types.SimpleNamespace(
select_graph=lambda name: graph,
)
module = types.SimpleNamespace(FalkorDB=lambda **kwargs: client)
monkeypatch.setitem(sys.modules, "falkordb", module)
return graph
@pytest.fixture()
def fake_neo4j(monkeypatch):
"""Install a fake `neo4j` module and return its recording graph handle."""
graph = _RecordingGraph()
session = _RecordingSession(graph)
driver = types.SimpleNamespace(
session=lambda: session,
close=lambda: None,
)
module = types.SimpleNamespace(GraphDatabase=types.SimpleNamespace(driver=lambda uri, auth=None: driver))
monkeypatch.setitem(sys.modules, "neo4j", module)
return graph
def test_falkordb_push_indexes_each_label_before_upserting(fake_falkordb):
"""Every MERGE label is indexed first, and before the first MERGE runs."""
from graphify.exporters.graphdb import push_to_falkordb
result = push_to_falkordb(_graph(), uri="localhost:6379", graph_name="graphify")
assert result == {"nodes": 4, "edges": 2}
indexes = [c for c in fake_falkordb.calls if c.startswith("CREATE INDEX")]
merges = [i for i, c in enumerate(fake_falkordb.calls) if c.startswith("MERGE")]
# One index per distinct label: the two code nodes share "Code", and
# "document" and "database" each get their own.
assert indexes == [
"CREATE INDEX FOR (n:Code) ON (n.id)",
"CREATE INDEX FOR (n:Database) ON (n.id)",
"CREATE INDEX FOR (n:Document) ON (n.id)",
]
assert merges, "expected the push to still upsert nodes"
assert fake_falkordb.calls.index(indexes[-1]) < merges[0]
def test_falkordb_push_indexes_before_the_upsert_loop_runs(fake_falkordb):
"""Ordering, stated as the property it protects: no MERGE precedes its index."""
from graphify.exporters.graphdb import push_to_falkordb
push_to_falkordb(_graph(), uri="localhost:6379", graph_name="graphify")
seen_index: set[str] = set()
for cypher in fake_falkordb.calls:
if cypher.startswith("CREATE INDEX FOR (n:"):
seen_index.add(cypher.split("(n:")[1].split(")")[0])
continue
if cypher.startswith("MERGE (n:"):
label = cypher.split("MERGE (n:")[1].split(" ")[0]
assert label in seen_index, f"MERGE on unindexed label {label}: {cypher}"
def test_falkordb_push_survives_an_existing_index(monkeypatch):
"""FalkorDB has no IF NOT EXISTS, so a re-push must not abort on the error."""
from graphify.exporters.graphdb import push_to_falkordb
graph = _RecordingGraph(fail_on_index=True)
client = types.SimpleNamespace(select_graph=lambda name: graph)
monkeypatch.setitem(sys.modules, "falkordb", types.SimpleNamespace(FalkorDB=lambda **kw: client))
result = push_to_falkordb(_graph(), uri="localhost:6379", graph_name="graphify")
# The push still upserts every node and edge despite the index errors.
assert result == {"nodes": 4, "edges": 2}
assert sum(1 for c in graph.calls if c.startswith("MERGE (n:")) == 4
assert sum(1 for c in graph.calls if c.startswith("CREATE INDEX")) == 3
def test_neo4j_push_indexes_each_label_before_upserting(fake_neo4j):
"""The Neo4j path had the same gap; it uses IF NOT EXISTS to stay idempotent."""
from graphify.exporters.graphdb import push_to_neo4j
result = push_to_neo4j(_graph(), uri="bolt://localhost:7687", user="neo4j", password="x")
assert result == {"nodes": 4, "edges": 2}
indexes = [c for c in fake_neo4j.calls if c.startswith("CREATE INDEX")]
merges = [i for i, c in enumerate(fake_neo4j.calls) if c.startswith("MERGE")]
assert indexes == [
"CREATE INDEX IF NOT EXISTS FOR (n:Code) ON (n.id)",
"CREATE INDEX IF NOT EXISTS FOR (n:Database) ON (n.id)",
"CREATE INDEX IF NOT EXISTS FOR (n:Document) ON (n.id)",
]
assert merges
assert fake_neo4j.calls.index(indexes[-1]) < merges[0]
def test_distinct_node_labels_sanitizes_like_the_push_loop():
"""The index set must not contain a label the MERGE loop would never write."""
from graphify.exporters.graphdb import _distinct_node_labels
G = nx.Graph()
# "code" -> "Code"; an all-punctuation value collapses to the "Entity"
# fallback, exactly as _safe_label does inside the push loops.
G.add_node("a", file_type="code")
G.add_node("b", file_type="Code")
G.add_node("c", file_type="!!!")
G.add_node("d") # no file_type at all -> "Entity"
assert _distinct_node_labels(G) == ["Code", "Entity"]