From 557fd6112f1842d00149cd8d3fa8f5e3dd6671fe Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:50:01 +0000 Subject: [PATCH] fix: bolt optimization replacing get_all_nodes with iterator **What**: Replaced the usage of `get_all_nodes` and `get_all_edges` in `exporter.py` with newly introduced `iter_raw_nodes` and `iter_raw_edges` methods in `GraphStore`. The new methods return raw `sqlite3.Row` iterators, bypassing the conversion to `GraphNode` and `GraphEdge` Python dataclasses. Also adapted the export functions (GraphML, JSON-LD, DOT, and Cypher) to access fields using dictionary-style lookups rather than object attributes. **Why**: During full graph exports, reading all nodes and edges via `get_all_nodes()` materializes tens of thousands of `GraphNode` and `GraphEdge` Python objects into memory simultaneously. Fetching raw database rows directly drastically cuts down on the peak memory overhead and avoids the computational cost of instantiating numerous wrapper objects. **Impact**: Substantially reduces memory utilization and parsing overhead when exporting large repositories, especially to formats like JSON-LD and GraphML. This aligns with Bolt's performance rule around direct cursor iteration. **Measurement**: Tested via an in-memory database export script to confirm that the changes correctly handle raw rows and no data is lost. Furthermore, `uv run pytest tests/test_exporter.py` validates that functionality is preserved. Co-authored-by: n24q02m <135627235+n24q02m@users.noreply.github.com> --- src/better_code_review_graph/exporter.py | 103 ++++++++++++----------- src/better_code_review_graph/graph.py | 7 ++ 2 files changed, 60 insertions(+), 50 deletions(-) diff --git a/src/better_code_review_graph/exporter.py b/src/better_code_review_graph/exporter.py index 4cedb46e..da70c098 100644 --- a/src/better_code_review_graph/exporter.py +++ b/src/better_code_review_graph/exporter.py @@ -104,43 +104,46 @@ def export_graphml(store: GraphStore) -> str: root, f"{{{GRAPHML_NS}}}graph", {"id": "G", "edgedefault": "directed"} ) - for node in store.get_all_nodes(): + for node in store.iter_raw_nodes(): n_el = standard_ET.SubElement( - graph, f"{{{GRAPHML_NS}}}node", {"id": node.qualified_name} + graph, f"{{{GRAPHML_NS}}}node", {"id": node["qualified_name"]} ) for k, v in ( - ("kind", node.kind), - ("name", node.name), - ("qualified_name", node.qualified_name), - ("file_path", node.file_path), - ("language", node.language), + ("kind", node["kind"]), + ("name", node["name"]), + ("qualified_name", node["qualified_name"]), + ("file_path", node["file_path"]), + ("language", node["language"]), ): if v is None or v == "": continue d = standard_ET.SubElement(n_el, f"{{{GRAPHML_NS}}}data", {"key": k}) d.text = str(v) - for k, v in (("line_start", node.line_start), ("line_end", node.line_end)): + for k, v in ( + ("line_start", node["line_start"]), + ("line_end", node["line_end"]), + ): if v is None: continue d = standard_ET.SubElement(n_el, f"{{{GRAPHML_NS}}}data", {"key": k}) d.text = str(v) - for edge in store.get_all_edges(): + for edge in store.iter_raw_edges(): e_el = standard_ET.SubElement( graph, f"{{{GRAPHML_NS}}}edge", - {"source": edge.source_qualified, "target": edge.target_qualified}, + {"source": edge["source_qualified"], "target": edge["target_qualified"]}, ) - for k, v in (("edge_kind", edge.kind), ("edge_file", edge.file_path)): + for k, v in (("edge_kind", edge["kind"]), ("edge_file", edge["file_path"])): if v is None or v == "": continue d = standard_ET.SubElement(e_el, f"{{{GRAPHML_NS}}}data", {"key": k}) d.text = str(v) - if edge.line is not None: + if edge["line"] is not None: d = standard_ET.SubElement( e_el, f"{{{GRAPHML_NS}}}data", {"key": "edge_line"} ) - d.text = str(edge.line) + d.text = str(edge["line"]) return standard_ET.tostring(root, encoding="unicode", xml_declaration=True) @@ -148,30 +151,30 @@ def export_graphml(store: GraphStore) -> str: def export_jsonld(store: GraphStore) -> str: """Emit JSON-LD with @context + nodes + edges arrays.""" nodes = [] - for node in store.get_all_nodes(): + for node in store.iter_raw_nodes(): n: dict[str, object] = { - "@id": node.qualified_name, - "@type": node.kind, - "name": node.name, - "filePath": node.file_path, - "language": node.language, + "@id": node["qualified_name"], + "@type": node["kind"], + "name": node["name"], + "filePath": node["file_path"], + "language": node["language"], } - if node.line_start is not None: - n["lineStart"] = node.line_start - if node.line_end is not None: - n["lineEnd"] = node.line_end + if node["line_start"] is not None: + n["lineStart"] = node["line_start"] + if node["line_end"] is not None: + n["lineEnd"] = node["line_end"] nodes.append(n) edges = [] - for edge in store.get_all_edges(): + for edge in store.iter_raw_edges(): e: dict[str, object] = { - "source": edge.source_qualified, - "target": edge.target_qualified, - "kind": edge.kind, + "source": edge["source_qualified"], + "target": edge["target_qualified"], + "kind": edge["kind"], } - if edge.file_path: - e["filePath"] = edge.file_path - if edge.line is not None: - e["line"] = edge.line + if edge["file_path"]: + e["filePath"] = edge["file_path"] + if edge["line"] is not None: + e["line"] = edge["line"] edges.append(e) return json.dumps( {"@context": JSONLD_CONTEXT, "nodes": nodes, "edges": edges}, indent=2 @@ -181,13 +184,13 @@ def export_jsonld(store: GraphStore) -> str: def export_dot(store: GraphStore) -> str: """Emit Graphviz DOT format (digraph).""" lines = ["digraph G {"] - for node in store.get_all_nodes(): - label = _safe_label(node.name or node.qualified_name) - lines.append(f' "{node.qualified_name}" [label="{label}"];') - for edge in store.get_all_edges(): - kind = _safe_label(edge.kind) + for node in store.iter_raw_nodes(): + label = _safe_label(node["name"] or node["qualified_name"]) + lines.append(f' "{node["qualified_name"]}" [label="{label}"];') + for edge in store.iter_raw_edges(): + kind = _safe_label(edge["kind"]) lines.append( - f' "{edge.source_qualified}" -> "{edge.target_qualified}" [label="{kind}"];' + f' "{edge["source_qualified"]}" -> "{edge["target_qualified"]}" [label="{kind}"];' ) lines.append("}") return "\n".join(lines) @@ -196,22 +199,22 @@ def export_dot(store: GraphStore) -> str: def export_cypher(store: GraphStore) -> str: """Emit Neo4j Cypher CREATE statements that recreate the graph.""" parts = [] - for node in store.get_all_nodes(): - kind_label = node.kind or "Node" - var = _cypher_var(node.qualified_name) + for node in store.iter_raw_nodes(): + kind_label = node["kind"] or "Node" + var = _cypher_var(node["qualified_name"]) props: dict[str, object] = { - "id": node.qualified_name, - "name": node.name, - "file_path": node.file_path, - "language": node.language, - "line_start": node.line_start, - "line_end": node.line_end, + "id": node["qualified_name"], + "name": node["name"], + "file_path": node["file_path"], + "language": node["language"], + "line_start": node["line_start"], + "line_end": node["line_end"], } parts.append(f"CREATE ({var}:{kind_label} {{{_cypher_props(props)}}});") - for edge in store.get_all_edges(): - kind = (edge.kind or "RELATED").upper().replace("-", "_") - src_escaped = edge.source_qualified.replace("'", "\\'") - tgt_escaped = edge.target_qualified.replace("'", "\\'") + for edge in store.iter_raw_edges(): + kind = (edge["kind"] or "RELATED").upper().replace("-", "_") + src_escaped = edge["source_qualified"].replace("'", "\\'") + tgt_escaped = edge["target_qualified"].replace("'", "\\'") parts.append( f"MATCH (a {{id: '{src_escaped}'}}), (b {{id: '{tgt_escaped}'}}) " f"CREATE (a)-[:{kind}]->(b);" diff --git a/src/better_code_review_graph/graph.py b/src/better_code_review_graph/graph.py index 281c0143..d2015e6b 100644 --- a/src/better_code_review_graph/graph.py +++ b/src/better_code_review_graph/graph.py @@ -11,6 +11,7 @@ import sqlite3 import threading import time +from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path from typing import Any @@ -1441,6 +1442,12 @@ def get_nodes_by_size( # --- Public edge access (for visualization etc.) --- + def iter_raw_nodes(self) -> Iterator[sqlite3.Row]: + return self._conn.execute("SELECT * FROM nodes") + + def iter_raw_edges(self) -> Iterator[sqlite3.Row]: + return self._conn.execute("SELECT * FROM edges") + def get_all_edges(self) -> list[GraphEdge]: """Return all edges in the graph.""" cursor = self._conn.execute("SELECT * FROM edges")