From a3275e2c970df7221ed85fb6eebed54984ec8133 Mon Sep 17 00:00:00 2001
From: Yi-111-a <153097222+Yi-111-a@users.noreply.github.com>
Date: Thu, 1 Oct 2026 03:11:51 +0800
Subject: [PATCH] fix(export): index (label, id) before the graph DB push
 upsert loop

`push_to_falkordb` and `push_to_neo4j` upsert every node with
`MERGE (n:<Label> {id: $id}) SET n += $props`, which matches on (label, id).
Neither path created an index on that pattern first, so the engine had to
scan every node already carrying the label to decide whether the node is
already there. Each upsert then costs O(nodes with that label) and a full push
degrades to O(n^2): on the graph in #3804 (167,428 nodes / 209,215 edges
across 6 labels) throughput collapsed from ~10,000 node upserts/60s to
~100 upserts/300s partway through, with the database at ~97% CPU.

Create one index per distinct label up front, before the upsert loop. The
label set is derived exactly as the push loops derive it, so every label a
MERGE matches against is covered. Neo4j uses IF NOT EXISTS to stay idempotent
across re-pushes; FalkorDB's CREATE INDEX has no such clause, so a re-push
against an already-indexed graph is tolerated rather than failing an
otherwise valid push.

Refs #3804

Co-Authored-By: opencode space-bunny-free <noreply@opencode.ai>
---
 graphify/exporters/graphdb.py      |  37 ++++++
 tests/test_graphdb_push_indexes.py | 180 +++++++++++++++++++++++++++++
 2 files changed, 217 insertions(+)
 create mode 100644 tests/test_graphdb_push_indexes.py

diff --git a/graphify/exporters/graphdb.py b/graphify/exporters/graphdb.py
index 14c47f0d5..d0bcc4f61 100644
--- a/graphify/exporters/graphdb.py
+++ b/graphify/exporters/graphdb.py
@@ -6,6 +6,19 @@ import networkx as nx
 import re
 
 
+def _distinct_node_labels(G: nx.Graph) -> list[str]:
+    """Distinct Cypher labels the node-upsert loop will write, sorted.
+
+    The label is derived exactly as the push loops derive it, so the indexes
+    created up front cover every label the MERGE statements will match against.
+    """
+    labels = set()
+    for _, data in G.nodes(data=True):
+        label = data.get("file_type", "Entity").capitalize()
+        labels.add(re.sub(r"[^A-Za-z0-9_]", "", label) or "Entity")
+    return sorted(labels)
+
+
 def push_to_neo4j(
     G: nx.Graph,
     uri: str,
@@ -42,6 +55,16 @@ def push_to_neo4j(
     edges_pushed = 0
 
     with driver.session() as session:
+        # Index (label, id) before any upsert, for the same reason as the
+        # FalkorDB path: MERGE matches on (label, id), so an unindexed match
+        # pattern scans the whole label and the push degrades to O(n^2)
+        # (#3804). IF NOT EXISTS keeps this idempotent across re-pushes.
+        for label in _distinct_node_labels(G):
+            try:
+                session.run(f"CREATE INDEX IF NOT EXISTS FOR (n:{label}) ON (n.id)")
+            except Exception:
+                pass
+
         for node_id, data in G.nodes(data=True):
             props = {
                 k: v for k, v in data.items()
@@ -141,6 +164,20 @@ def push_to_falkordb(
     nodes_pushed = 0
     edges_pushed = 0
 
+    # Index (label, id) before any upsert. The MERGE below matches on
+    # (label, id), so with no index FalkorDB scans every existing node carrying
+    # that label to decide whether the node is new: each upsert costs
+    # O(nodes with that label) and the whole push degrades to O(n^2) (#3804).
+    # Measured on a 167k-node / 209k-edge graph, adding these indexes first
+    # restored the throughput the push started at. FalkorDB's CREATE INDEX has
+    # no IF NOT EXISTS, so a re-run against an already-indexed graph raises —
+    # tolerate that rather than failing an otherwise valid push.
+    for label in _distinct_node_labels(G):
+        try:
+            graph.query(f"CREATE INDEX FOR (n:{label}) ON (n.id)")
+        except Exception:
+            pass
+
     for node_id, data in G.nodes(data=True):
         props = {
             k: v for k, v in data.items()
diff --git a/tests/test_graphdb_push_indexes.py b/tests/test_graphdb_push_indexes.py
new file mode 100644
index 000000000..15fe0605a
--- /dev/null
+++ b/tests/test_graphdb_push_indexes.py
@@ -0,0 +1,180 @@
+"""#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"]