From bf5eebf7bcf127913750f884f0145f6f454aa8bd Mon Sep 17 00:00:00 2001 From: BingqingLyu Date: Fri, 22 May 2026 16:02:41 +0800 Subject: [PATCH 01/10] feat: add NeuG as parallel graph storage engine with Cypher query support Integrate NeuG as an optional storage backend alongside NetworkX. During extract, graph data is written to both graph.json (NetworkX) and graph.db (NeuG) when neug is installed. Adds `graphify cypher` CLI command and `cypher_query` MCP tool for direct Cypher queries. - New graphify/storage.py: NeuG adapter (init, ingest, query, close) - __main__.py: NeuG write in extract flow + `cypher` CLI command - serve.py: NeuG connection init + `cypher_query` MCP tool - pyproject.toml: neug optional dependency - Tests: unit tests, CLI tests, E2E integration script Co-Authored-By: Claude Opus 4.6 --- graphify/__main__.py | 54 +++++ graphify/serve.py | 45 +++++ graphify/storage.py | 295 +++++++++++++++++++++++++++ pyproject.toml | 3 +- tests/test_cypher_cli.py | 58 ++++++ tests/test_neug_e2e.sh | 418 +++++++++++++++++++++++++++++++++++++++ tests/test_storage.py | 177 +++++++++++++++++ 7 files changed, 1049 insertions(+), 1 deletion(-) create mode 100644 graphify/storage.py create mode 100644 tests/test_cypher_cli.py create mode 100755 tests/test_neug_e2e.sh create mode 100644 tests/test_storage.py diff --git a/graphify/__main__.py b/graphify/__main__.py index a2678655e7..2d6b769027 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -1469,6 +1469,8 @@ def main() -> None: print(" cluster-only rerun clustering on an existing graph.json and regenerate report") print(" --no-viz skip graph.html generation (useful for >5000 node graphs / CI)") print(" --graph path to graph.json (default /graphify-out/graph.json)") + print(" cypher \"MATCH ...\" execute a Cypher query against graph.db (requires neug)") + print(" --db path to graph.db (default graphify-out/graph.db)") print(" query \"\" BFS traversal of graph.json for a question") print(" --dfs use depth-first instead of breadth-first") print(" --context C explicit edge-context filter (repeatable)") @@ -1787,6 +1789,31 @@ def main() -> None: else: print("Usage: graphify hook [install|uninstall|status]", file=sys.stderr) sys.exit(1) + elif cmd == "cypher": + if len(sys.argv) < 3: + print('Usage: graphify cypher "MATCH ..." [--db path]', file=sys.stderr) + sys.exit(1) + query_str = sys.argv[2] + db_path = str(Path(_GRAPHIFY_OUT) / "graph.db") + args = sys.argv[3:] + for i, a in enumerate(args): + if a == "--db" and i + 1 < len(args): + db_path = args[i + 1] + try: + from graphify.storage import init_db, execute_cypher, close_db + except ImportError: + print("error: neug is not installed. Run: pip install neug", file=sys.stderr) + sys.exit(1) + if not Path(db_path).exists(): + print(f"error: database not found: {db_path}", file=sys.stderr) + sys.exit(1) + db, conn = init_db(db_path) + try: + results = execute_cypher(conn, query_str) + for row in results: + print("\t".join(str(v) for v in row)) + finally: + close_db(db, conn) elif cmd == "query": if len(sys.argv) < 3: print("Usage: graphify query \"\" [--dfs] [--context C] [--budget N] [--graph path]", file=sys.stderr) @@ -3314,6 +3341,19 @@ def _progress(idx: int, total: int, _result: dict) -> None: graph_json_path.write_text( json.dumps(merged, indent=2), encoding="utf-8" ) + try: + from graphify.storage import init_db as _init_db, ingest_extraction as _ingest, close_db as _close_db + _db_path = str(graphify_out / "graph.db") + _is_inc = Path(_db_path).exists() + _db, _conn = _init_db(_db_path) + _ingest(_conn, merged, incremental=_is_inc, + prune_sources=deleted_files or None, root=target) + _close_db(_db, _conn) + print("[graphify extract] graph.db written (powered by NeuG)") + except ImportError: + pass + except Exception as _exc: + print(f"[graphify extract] warning: NeuG write failed: {_exc}", file=sys.stderr) cost = _estimate_cost( backend, merged["input_tokens"], merged["output_tokens"] ) @@ -3391,6 +3431,20 @@ def _progress(idx: int, total: int, _result: dict) -> None: from graphify.export import backup_if_protected as _backup _backup(graphify_out) _to_json(G, communities, str(graph_json_path), force=True) + try: + from graphify.storage import init_db as _init_db, ingest_extraction as _ingest, ingest_communities as _ingest_comm, close_db as _close_db + _db_path = str(graphify_out / "graph.db") + _is_inc = Path(_db_path).exists() + _db, _conn = _init_db(_db_path) + _ingest(_conn, merged, incremental=_is_inc, + prune_sources=deleted_files or None, root=target) + _ingest_comm(_conn, communities) + _close_db(_db, _conn) + print("[graphify extract] graph.db written (powered by NeuG)") + except ImportError: + pass + except Exception as _exc: + print(f"[graphify extract] warning: NeuG write failed: {_exc}", file=sys.stderr) if merged.get("output_tokens", 0) > 0: (graphify_out / ".graphify_semantic_marker").write_text( json.dumps({"output_tokens": merged["output_tokens"]}), encoding="utf-8" diff --git a/graphify/serve.py b/graphify/serve.py index 6e5d4a1f63..fae867081f 100644 --- a/graphify/serve.py +++ b/graphify/serve.py @@ -487,6 +487,20 @@ def serve(graph_path: str = "graphify-out/graph.json") -> None: G = _load_graph(graph_path) communities = _communities_from_graph(G) + _neug_conn = None + _neug_db = None + _neug_execute = None + try: + from graphify.storage import init_db as _neug_init, execute_cypher as _neug_exec, close_db as _neug_close + _neug_db_path = str(Path(graph_path).parent / "graph.db") + if Path(_neug_db_path).exists(): + _neug_db, _neug_conn = _neug_init(_neug_db_path) + _neug_execute = _neug_exec + except ImportError: + pass + except Exception: + pass + # Hot-reload state: mtime+size key lets us detect graph.json changes without # polling. Initialised from the file stat at startup so the first tool call # never triggers a redundant reload. @@ -646,6 +660,20 @@ async def list_tools() -> list[types.Tool]: }, }, ), + types.Tool( + name="cypher_query", + description=( + "Execute a Cypher query against the NeuG graph database. " + "Returns tabular results. Requires neug to be installed and graph.db to exist." + ), + inputSchema={ + "type": "object", + "properties": { + "query": {"type": "string", "description": "Cypher query string"}, + }, + "required": ["query"], + }, + ), ] def _tool_query_graph(arguments: dict) -> str: @@ -882,6 +910,22 @@ def _tool_triage_prs(arguments: dict) -> str: ) return "\n\n".join(lines) + def _tool_cypher_query(arguments: dict) -> str: + if _neug_conn is None: + return "NeuG not available (not installed or graph.db not found)." + query = arguments["query"] + from graphify.storage import execute_cypher as _exec_cypher + try: + results = _exec_cypher(_neug_conn, query) + except RuntimeError as exc: + return f"Cypher error: {exc}" + if not results: + return "No results." + lines = [] + for row in results: + lines.append("\t".join(str(v) for v in row)) + return "\n".join(lines) + _handlers = { "query_graph": _tool_query_graph, "get_node": _tool_get_node, @@ -893,6 +937,7 @@ def _tool_triage_prs(arguments: dict) -> str: "list_prs": _tool_list_prs, "get_pr_impact": _tool_get_pr_impact, "triage_prs": _tool_triage_prs, + "cypher_query": _tool_cypher_query, } def _load_community_labels() -> dict[int, str]: diff --git a/graphify/storage.py b/graphify/storage.py new file mode 100644 index 0000000000..3c15c3fc2e --- /dev/null +++ b/graphify/storage.py @@ -0,0 +1,295 @@ +"""NeuG graph database adapter for graphify. + +Provides an optional parallel storage engine alongside NetworkX. +All functions are guarded by `import neug` — when NeuG is not installed, +callers should catch ImportError and skip silently. +""" +from __future__ import annotations + +import os +import re +import unicodedata +from pathlib import Path + +import neug + +from .build import _FILE_TYPE_SYNONYMS, _normalize_id, _norm_source_file +from .validate import VALID_FILE_TYPES + +# --------------------------------------------------------------------------- +# Node tables (one per file_type) +# --------------------------------------------------------------------------- + +_NODE_TABLES = { + "code": """CREATE NODE TABLE IF NOT EXISTS code ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, source_location STRING, community INT64)""", + "document": """CREATE NODE TABLE IF NOT EXISTS document ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, community INT64)""", + "paper": """CREATE NODE TABLE IF NOT EXISTS paper ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, community INT64)""", + "image": """CREATE NODE TABLE IF NOT EXISTS image ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, community INT64)""", + "concept": """CREATE NODE TABLE IF NOT EXISTS concept ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, community INT64)""", + "rationale": """CREATE NODE TABLE IF NOT EXISTS rationale ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, community INT64)""", +} + +# --------------------------------------------------------------------------- +# Pre-built edge tables (based on actual graphify data model) +# --------------------------------------------------------------------------- + +_EDGE_PAIRS: list[tuple[str, str]] = [ + ("code", "code"), + ("code", "rationale"), + ("code", "concept"), + ("code", "document"), + ("document", "document"), + ("document", "concept"), + ("paper", "concept"), + ("paper", "paper"), + ("concept", "concept"), + ("image", "concept"), +] + +_EDGE_DDL_TEMPLATE = """CREATE REL TABLE IF NOT EXISTS edge_{src}_{tgt}( + FROM {src} TO {tgt}, + relation STRING, confidence STRING, + confidence_score DOUBLE, source_file STRING, weight DOUBLE)""" + +_created_rel_tables: set[tuple[str, str]] = set() + + +def _edge_table_name(src_type: str, tgt_type: str) -> str: + return f"edge_{src_type}_{tgt_type}" + + +# --------------------------------------------------------------------------- +# Cypher string escaping +# --------------------------------------------------------------------------- + + +def _cesc(value: str) -> str: + return ( + value + .replace("\\", "\\\\") + .replace("'", "\\'") + .replace("\n", "\\n") + .replace("\r", "\\r") + .replace("\t", "\\t") + ) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def init_db(db_path: str) -> tuple[neug.Database, object]: + """Create or open a NeuG database and initialize schema.""" + db = neug.Database(db_path) + conn = db.connect() + + for ddl in _NODE_TABLES.values(): + conn.execute(ddl) + + _created_rel_tables.clear() + for src, tgt in _EDGE_PAIRS: + conn.execute(_EDGE_DDL_TEMPLATE.format(src=src, tgt=tgt)) + _created_rel_tables.add((src, tgt)) + + return db, conn + + +def _ensure_rel_table(conn: object, src_type: str, tgt_type: str) -> None: + """Create an edge table on-the-fly if not already pre-built.""" + pair = (src_type, tgt_type) + if pair in _created_rel_tables: + return + conn.execute(_EDGE_DDL_TEMPLATE.format(src=src_type, tgt=tgt_type)) + _created_rel_tables.add(pair) + + +def _fix_file_type(ft: str | None) -> str: + """Canonicalize file_type, matching build.py:138-146 logic.""" + if not ft or ft not in VALID_FILE_TYPES: + return _FILE_TYPE_SYNONYMS.get(ft, "concept") if ft else "concept" + return ft + + +def ingest_extraction( + conn: object, + extraction: dict, + *, + incremental: bool = False, + prune_sources: list[str] | None = None, + root: str | Path | None = None, +) -> None: + """Write an extraction dict into NeuG. + + incremental=False: first build — uses CREATE (faster). + incremental=True: update — uses MERGE (upsert). + """ + _root = str(Path(root).resolve()) if root else None + + # --- prune deleted/changed files first --- + if prune_sources: + for sf in prune_sources: + sf_norm = _norm_source_file(sf, _root) or sf + for tbl in _NODE_TABLES: + conn.execute( + f"MATCH (n:{tbl}) WHERE n.source_file = '{_cesc(sf_norm)}' " + f"DETACH DELETE n" + ) + + # --- build node lookup: id -> file_type --- + node_types: dict[str, str] = {} + nodes = extraction.get("nodes") or [] + edges = extraction.get("edges") or [] + + # --- write nodes --- + _written_ids: set[str] = set() + _n_errors = 0 + for node in nodes: + nid = _normalize_id(node.get("id", "")) + if not nid: + continue + ft = _fix_file_type(node.get("file_type")) + label = node.get("label", "") + sf = _norm_source_file(node.get("source_file"), _root) or "" + sl = node.get("source_location") or "" + node_types[nid] = ft + if nid in _written_ids: + continue + _written_ids.add(nid) + + try: + if incremental: + existing = list(conn.execute( + f"MATCH (n:{ft} {{id: '{_cesc(nid)}'}}) RETURN n.id" + )) + if existing: + props = f"n.label = '{_cesc(label)}', n.source_file = '{_cesc(sf)}'" + if ft == "code": + props += f", n.source_location = '{_cesc(sl)}'" + conn.execute( + f"MATCH (n:{ft} {{id: '{_cesc(nid)}'}}) SET {props}" + ) + else: + if ft == "code": + conn.execute( + f"CREATE (n:code {{id: '{_cesc(nid)}', " + f"label: '{_cesc(label)}', " + f"source_file: '{_cesc(sf)}', " + f"source_location: '{_cesc(sl)}'}})" + ) + else: + conn.execute( + f"CREATE (n:{ft} {{id: '{_cesc(nid)}', " + f"label: '{_cesc(label)}', " + f"source_file: '{_cesc(sf)}'}})" + ) + else: + if ft == "code": + conn.execute( + f"CREATE (n:code {{id: '{_cesc(nid)}', " + f"label: '{_cesc(label)}', " + f"source_file: '{_cesc(sf)}', " + f"source_location: '{_cesc(sl)}'}})" + ) + else: + conn.execute( + f"CREATE (n:{ft} {{id: '{_cesc(nid)}', " + f"label: '{_cesc(label)}', " + f"source_file: '{_cesc(sf)}'}})" + ) + except RuntimeError: + _n_errors += 1 + + # --- write edges --- + _e_errors = 0 + for edge in edges: + src_key = edge.get("source") or edge.get("from", "") + tgt_key = edge.get("target") or edge.get("to", "") + src_id = _normalize_id(src_key) + tgt_id = _normalize_id(tgt_key) + if not src_id or not tgt_id: + continue + + src_ft = node_types.get(src_id) + tgt_ft = node_types.get(tgt_id) + if not src_ft or not tgt_ft: + continue + + _ensure_rel_table(conn, src_ft, tgt_ft) + tbl = _edge_table_name(src_ft, tgt_ft) + rel = _cesc(edge.get("relation", "")) + conf = _cesc(edge.get("confidence", "")) + conf_score = float(edge.get("confidence_score", 0.0)) + e_sf = _cesc(_norm_source_file(edge.get("source_file"), _root) or "") + weight = float(edge.get("weight", 1.0)) + + try: + if incremental: + conn.execute( + f"MATCH (a:{src_ft} {{id: '{_cesc(src_id)}'}}), " + f"(b:{tgt_ft} {{id: '{_cesc(tgt_id)}'}}) " + f"CREATE (a)-[:{tbl} {{relation: '{rel}', confidence: '{conf}', " + f"confidence_score: {conf_score}, source_file: '{e_sf}', " + f"weight: {weight}}}]->(b)" + ) + else: + conn.execute( + f"MATCH (a:{src_ft} {{id: '{_cesc(src_id)}'}}), " + f"(b:{tgt_ft} {{id: '{_cesc(tgt_id)}'}}) " + f"CREATE (a)-[:{tbl} {{relation: '{rel}', confidence: '{conf}', " + f"confidence_score: {conf_score}, source_file: '{e_sf}', " + f"weight: {weight}}}]->(b)" + ) + except RuntimeError: + _e_errors += 1 + + if _n_errors or _e_errors: + import logging + logging.getLogger(__name__).warning( + "NeuG ingest: %d node(s) and %d edge(s) skipped due to errors", + _n_errors, _e_errors, + ) + + +def ingest_communities( + conn: object, + communities: dict[int, list[str]], + community_labels: dict[int, str] | None = None, +) -> None: + """Write community assignments into NeuG node properties.""" + for cid, node_ids in communities.items(): + for nid in node_ids: + nid_norm = _normalize_id(nid) + if not nid_norm: + continue + for tbl in _NODE_TABLES: + conn.execute( + f"MATCH (n:{tbl}) WHERE n.id = '{_cesc(nid_norm)}' " + f"SET n.community = {int(cid)}" + ) + + +def execute_cypher(conn: object, query: str) -> list[list]: + """Execute a Cypher query and return results as list of lists.""" + try: + return list(conn.execute(query)) + except RuntimeError as exc: + raise RuntimeError(f"Cypher query failed: {exc}") from exc + + +def close_db(db: neug.Database, conn: object) -> None: + """Close the NeuG connection and database.""" + conn.close() + db.close() diff --git a/pyproject.toml b/pyproject.toml index 57b595babf..b160f8ddf0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,7 +65,8 @@ gemini = ["openai", "tiktoken"] openai = ["openai", "tiktoken"] chinese = ["jieba"] sql = ["tree-sitter-sql"] -all = ["mcp", "neo4j", "pypdf", "markdownify", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper", "yt-dlp", "matplotlib", "openai", "tiktoken", "boto3", "tree-sitter-sql", "jieba"] +neug = ["neug"] +all = ["mcp", "neo4j", "neug", "pypdf", "markdownify", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper", "yt-dlp", "matplotlib", "openai", "tiktoken", "boto3", "tree-sitter-sql", "jieba"] [project.scripts] graphify = "graphify.__main__:main" diff --git a/tests/test_cypher_cli.py b/tests/test_cypher_cli.py new file mode 100644 index 0000000000..fb705d9b5d --- /dev/null +++ b/tests/test_cypher_cli.py @@ -0,0 +1,58 @@ +"""Tests for the `graphify cypher` CLI command.""" +import json +import subprocess +import sys +import tempfile +from pathlib import Path + +import pytest + +try: + import neug + _has_neug = True +except ImportError: + _has_neug = False + +pytestmark = pytest.mark.skipif(not _has_neug, reason="neug not installed") + +FIXTURES = Path(__file__).parent / "fixtures" +EXTRACTION_JSON = FIXTURES / "extraction.json" + + +def _build_db(tmp_path) -> str: + from graphify.storage import init_db, ingest_extraction, close_db + db_path = str(tmp_path / "graph.db") + ext = json.loads(EXTRACTION_JSON.read_text()) + db, conn = init_db(db_path) + ingest_extraction(conn, ext, incremental=False) + close_db(db, conn) + return db_path + + +def test_cypher_command_basic(tmp_path): + db_path = _build_db(tmp_path) + result = subprocess.run( + [sys.executable, "-m", "graphify", "cypher", + "MATCH (n:code) RETURN count(n)", "--db", db_path], + capture_output=True, text=True, timeout=30, + ) + assert result.returncode == 0 + assert "3" in result.stdout + + +def test_cypher_command_db_not_found(tmp_path): + result = subprocess.run( + [sys.executable, "-m", "graphify", "cypher", + "MATCH (n) RETURN n", "--db", str(tmp_path / "nonexistent.db")], + capture_output=True, text=True, timeout=30, + ) + assert result.returncode != 0 + assert "not found" in result.stderr.lower() or "error" in result.stderr.lower() + + +def test_cypher_command_no_query(): + result = subprocess.run( + [sys.executable, "-m", "graphify", "cypher"], + capture_output=True, text=True, timeout=30, + ) + assert result.returncode != 0 diff --git a/tests/test_neug_e2e.sh b/tests/test_neug_e2e.sh new file mode 100755 index 0000000000..48b5a246f8 --- /dev/null +++ b/tests/test_neug_e2e.sh @@ -0,0 +1,418 @@ +#!/usr/bin/env bash +# +# End-to-end test for graphify + NeuG integration. +# +# Tests the full flow: +# 1. graphify extract (AST-only, no LLM) → graph.json + graph.db +# 2. graphify cypher → query against graph.db +# 3. Incremental re-extract → graph.db updated +# 4. MCP server tool registration (smoke test) +# +# Prerequisites: +# - pip install neug +# - graphify installed from current source (pip install -e .) +# +# Usage: +# bash tests/test_neug_e2e.sh +# +set -uo pipefail + +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +NC='\033[0m' + +PASS=0 +FAIL=0 +SKIP=0 + +pass() { echo -e " ${GREEN}PASS${NC}: $1"; ((PASS++)); } +fail() { echo -e " ${RED}FAIL${NC}: $1"; ((FAIL++)); } +skip() { echo -e " ${YELLOW}SKIP${NC}: $1"; ((SKIP++)); } + +echo "======================================" +echo " graphify + NeuG E2E Integration Test" +echo "======================================" +echo "" + +# --- Check prerequisites --- +if ! python3 -c "import neug" 2>/dev/null; then + echo "ERROR: neug not installed. Run: pip install neug" + exit 1 +fi + +if ! python3 -c "import graphify" 2>/dev/null; then + echo "ERROR: graphify not importable. Run: pip install -e . from graphify root" + exit 1 +fi + +# --- Setup test project --- +TMPDIR=$(mktemp -d) +trap "rm -rf $TMPDIR" EXIT + +PROJECT="$TMPDIR/sample_project" +mkdir -p "$PROJECT/src" + +cat > "$PROJECT/src/main.py" << 'PYEOF' +"""Main application module.""" + +class Database: + """Database connection manager.""" + def __init__(self, url: str): + self.url = url + self.conn = None + + def connect(self): + """Establish connection.""" + from src.utils import validate_url + validate_url(self.url) + self.conn = True + return self + + def query(self, sql: str): + """Execute a query.""" + if not self.conn: + raise RuntimeError("Not connected") + return [] + + +class App: + """Main application.""" + def __init__(self): + self.db = Database("localhost:5432") + + def run(self): + self.db.connect() + results = self.db.query("SELECT 1") + return results +PYEOF + +cat > "$PROJECT/src/utils.py" << 'PYEOF' +"""Utility functions.""" + +def validate_url(url: str) -> bool: + """Validate a database URL.""" + if not url: + raise ValueError("Empty URL") + return ":" in url + + +def format_result(row: dict) -> str: + """Format a query result row.""" + return ", ".join(f"{k}={v}" for k, v in row.items()) + + +class Logger: + """Simple logger.""" + def __init__(self, name: str): + self.name = name + + def info(self, msg: str): + print(f"[{self.name}] {msg}") +PYEOF + +cat > "$PROJECT/src/models.py" << 'PYEOF' +"""Data models.""" +from src.utils import Logger + +class User: + """User model.""" + def __init__(self, name: str, email: str): + self.name = name + self.email = email + self.logger = Logger("User") + + def save(self): + self.logger.info(f"Saving user {self.name}") + +class Session: + """Session model.""" + def __init__(self, user: User): + self.user = user + self.active = True + + def close(self): + self.active = False +PYEOF + +echo "Test project created at $PROJECT" +echo "" + +# ============================================================ +# TEST 1: First extract (AST-only, no cluster) +# ============================================================ +echo "[Test 1] graphify extract (first build, no-semantic, no-cluster)" + +OUTPUT=$(cd "$PROJECT" && GEMINI_API_KEY=dummy python3 -m graphify extract . --no-semantic --no-cluster 2>&1) + +CLEAN_OUTPUT=$(echo "$OUTPUT" | grep -v "^INFO\|^E20") +if echo "$CLEAN_OUTPUT" | grep -q "graph.db written"; then + pass "graph.db written message present" +else + if [ -d "$PROJECT/graphify-out/graph.db" ]; then + pass "graph.db created (message may be suppressed)" + else + fail "graph.db NOT created" + echo " Output: $(echo "$CLEAN_OUTPUT" | grep -i 'neug\|graph.db\|warning\|error' | head -5)" + fi +fi + +if [ -f "$PROJECT/graphify-out/graph.json" ]; then + pass "graph.json created" +else + fail "graph.json NOT created" +fi + +# Verify graph.json has nodes +NODE_COUNT=$(python3 -c " +import json +d = json.load(open('$PROJECT/graphify-out/graph.json')) +print(len(d.get('nodes', []))) +") +if [ "$NODE_COUNT" -gt 0 ]; then + pass "graph.json has $NODE_COUNT nodes" +else + fail "graph.json has 0 nodes" +fi + +echo "" + +# ============================================================ +# TEST 2: graphify cypher — query the database +# ============================================================ +echo "[Test 2] graphify cypher — Cypher queries against graph.db" + +DB_PATH="$PROJECT/graphify-out/graph.db" + +if [ ! -e "$DB_PATH" ]; then + skip "graph.db not available, skipping cypher tests" +else + # Count all nodes + CYPHER_COUNT=$(python3 -m graphify cypher "MATCH (n:code) RETURN count(n)" --db "$DB_PATH" 2>/dev/null | grep -v "^INFO\|^E20") + if [ -n "$CYPHER_COUNT" ] && [ "$CYPHER_COUNT" != "0" ]; then + pass "cypher count query returned: $CYPHER_COUNT" + else + fail "cypher count query returned empty or zero" + fi + + # Query node labels + CYPHER_LABELS=$(python3 -m graphify cypher "MATCH (n:code) RETURN n.label LIMIT 5" --db "$DB_PATH" 2>/dev/null | grep -v "^INFO\|^E20") + if [ -n "$CYPHER_LABELS" ]; then + pass "cypher label query returned results" + else + fail "cypher label query returned empty" + fi + + # Query edges + CYPHER_EDGES=$(python3 -m graphify cypher "MATCH (a:code)-[e:edge_code_code]->(b:code) RETURN a.label, e.relation, b.label LIMIT 5" --db "$DB_PATH" 2>/dev/null | grep -v "^INFO\|^E20") + if [ -n "$CYPHER_EDGES" ]; then + pass "cypher edge query returned results" + else + fail "cypher edge query returned empty (may have no code->code edges)" + fi + + # Error case: bad query + BAD_RESULT=$(python3 -m graphify cypher "INVALID CYPHER" --db "$DB_PATH" 2>&1 | grep -v "^INFO\|^E20" || true) + if echo "$BAD_RESULT" | grep -qi "error\|fail\|traceback"; then + pass "bad cypher query properly errors" + else + fail "bad cypher query did not error" + fi + + # Error case: missing db + MISS_RESULT=$(python3 -m graphify cypher "MATCH (n) RETURN n" --db "/nonexistent/path.db" 2>&1 | grep -v "^INFO\|^E20" || true) + if echo "$MISS_RESULT" | grep -qi "not found\|error"; then + pass "missing db properly errors" + else + fail "missing db did not error" + fi +fi + +echo "" + +# ============================================================ +# TEST 3: Incremental extract (modify a file, re-run) +# ============================================================ +echo "[Test 3] Incremental extract (modify file, re-run)" + +# Modify a file +cat >> "$PROJECT/src/utils.py" << 'PYEOF' + +def new_function(): + """A newly added function.""" + return 42 +PYEOF + +OUTPUT2=$(cd "$PROJECT" && GEMINI_API_KEY=dummy python3 -m graphify extract . --no-semantic --no-cluster 2>&1) + +CLEAN_OUTPUT2=$(echo "$OUTPUT2" | grep -v "^INFO\|^E20") +if echo "$CLEAN_OUTPUT2" | grep -q "incremental"; then + pass "incremental mode detected" +else + skip "incremental mode not detected in output" +fi + +# Check that graph.db still exists and is queryable +if [ -e "$DB_PATH" ]; then + CYPHER_INC=$(python3 -m graphify cypher "MATCH (n:code) RETURN count(n)" --db "$DB_PATH" 2>/dev/null | grep -v "^INFO\|^E20" || echo "ERROR") + if [ "$CYPHER_INC" != "ERROR" ] && [ -n "$CYPHER_INC" ]; then + pass "graph.db still queryable after incremental ($CYPHER_INC nodes)" + else + fail "graph.db not queryable after incremental" + fi +else + skip "graph.db not available for incremental test" +fi + +echo "" + +# ============================================================ +# TEST 4: Python API — storage module direct test +# ============================================================ +echo "[Test 4] Python API — storage module" + +python3 << 'PYTEST' +import sys, json, tempfile, os +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath("__file__")))) +from graphify.storage import init_db, ingest_extraction, ingest_communities, execute_cypher, close_db + +# Small extraction (under 4096 limit) +extraction = { + "nodes": [ + {"id": "fn_main", "label": "main()", "file_type": "code", "source_file": "app.py", "source_location": "L1"}, + {"id": "fn_helper", "label": "helper()", "file_type": "code", "source_file": "app.py", "source_location": "L10"}, + {"id": "concept_arch", "label": "architecture", "file_type": "concept", "source_file": "docs.md"}, + ], + "edges": [ + {"source": "fn_main", "target": "fn_helper", "relation": "calls", "confidence": "EXTRACTED", "source_file": "app.py", "weight": 1.0}, + {"source": "fn_main", "target": "concept_arch", "relation": "implements", "confidence": "INFERRED", "source_file": "app.py", "weight": 0.8}, + ], +} + +d = tempfile.mkdtemp() +db_path = os.path.join(d, "test.db") +db, conn = init_db(db_path) + +# First build +ingest_extraction(conn, extraction, incremental=False) +nodes = execute_cypher(conn, "MATCH (n:code) RETURN count(n)") +assert nodes[0][0] == 2, f"Expected 2 code nodes, got {nodes[0][0]}" + +edges = execute_cypher(conn, "MATCH ()-[e:edge_code_code]->() RETURN count(e)") +assert edges[0][0] == 1, f"Expected 1 code-code edge, got {edges[0][0]}" + +# Communities +ingest_communities(conn, {0: ["fn_main", "fn_helper"], 1: ["concept_arch"]}) +comm = execute_cypher(conn, "MATCH (n:code {id: 'fn_main'}) RETURN n.community") +assert comm[0][0] == 0, f"Expected community 0, got {comm[0][0]}" + +# Incremental update +extraction["nodes"][0]["label"] = "main_v2()" +ingest_extraction(conn, extraction, incremental=True) +updated = execute_cypher(conn, "MATCH (n:code {id: 'fn_main'}) RETURN n.label") +assert updated[0][0] == "main_v2()", f"Expected 'main_v2()', got {updated[0][0]}" + +close_db(db, conn) +print(" All Python API assertions passed") +PYTEST + +if [ $? -eq 0 ]; then + pass "Python API test passed" +else + fail "Python API test failed" +fi + +echo "" + +# ============================================================ +# TEST 5: MCP server tool registration (smoke test) +# ============================================================ +echo "[Test 5] MCP server — cypher_query tool registered" + +python3 << 'MCPTEST' +import sys +# Check that serve.py has cypher_query in its tool list +import importlib.util +spec = importlib.util.find_spec("graphify.serve") +if spec is None: + print(" graphify.serve not found") + sys.exit(1) + +source = open(spec.origin).read() +if "cypher_query" in source and "_tool_cypher_query" in source: + print(" cypher_query tool found in serve.py") + sys.exit(0) +else: + print(" cypher_query tool NOT found in serve.py") + sys.exit(1) +MCPTEST + +if [ $? -eq 0 ]; then + pass "cypher_query tool registered in MCP server" +else + fail "cypher_query tool not found in MCP server" +fi + +echo "" + +# ============================================================ +# TEST 6: Graceful fallback when neug not installed +# ============================================================ +echo "[Test 6] Graceful fallback (simulated)" + +python3 << 'FALLBACK' +import sys, importlib, types + +# Simulate neug not being importable by temporarily removing it +saved = sys.modules.pop("neug", None) +saved_storage = sys.modules.pop("graphify.storage", None) + +# Create a fake module that raises ImportError +blocker = types.ModuleType("neug") +blocker.__spec__ = None + +class NeuGBlocker: + def find_module(self, name, path=None): + if name == "neug" or name.startswith("neug."): + return self + def load_module(self, name): + raise ImportError("simulated: neug not installed") + +sys.meta_path.insert(0, NeuGBlocker()) + +try: + # This should raise ImportError (caught by __main__.py) + from graphify.storage import init_db + print(" ERROR: import should have failed") + sys.exit(1) +except ImportError: + print(" ImportError correctly raised when neug missing") + sys.exit(0) +finally: + sys.meta_path.pop(0) + if saved: + sys.modules["neug"] = saved + if saved_storage: + sys.modules["graphify.storage"] = saved_storage +FALLBACK + +if [ $? -eq 0 ]; then + pass "graceful fallback when neug not installed" +else + fail "fallback test failed" +fi + +echo "" + +# ============================================================ +# Summary +# ============================================================ +echo "======================================" +TOTAL=$((PASS + FAIL + SKIP)) +echo -e " Results: ${GREEN}$PASS passed${NC}, ${RED}$FAIL failed${NC}, ${YELLOW}$SKIP skipped${NC} / $TOTAL total" +echo "======================================" + +if [ $FAIL -gt 0 ]; then + exit 1 +fi +exit 0 diff --git a/tests/test_storage.py b/tests/test_storage.py new file mode 100644 index 0000000000..839d9c9abd --- /dev/null +++ b/tests/test_storage.py @@ -0,0 +1,177 @@ +"""Tests for graphify.storage — NeuG adapter layer.""" +import json +import shutil +import tempfile +from pathlib import Path + +import pytest + +try: + import neug + _has_neug = True +except ImportError: + _has_neug = False + +pytestmark = pytest.mark.skipif(not _has_neug, reason="neug not installed") + +FIXTURES = Path(__file__).parent / "fixtures" +EXTRACTION_JSON = FIXTURES / "extraction.json" + + +def _load_extraction() -> dict: + return json.loads(EXTRACTION_JSON.read_text()) + + +@pytest.fixture() +def tmp_db(tmp_path): + db_path = str(tmp_path / "test.db") + yield db_path + + +def _init(db_path): + from graphify.storage import init_db + return init_db(db_path) + + +def _close(db, conn): + from graphify.storage import close_db + close_db(db, conn) + + +def _query(conn, cypher): + from graphify.storage import execute_cypher + return execute_cypher(conn, cypher) + + +# --- init_db --- + +def test_init_db_creates_tables(tmp_db): + db, conn = _init(tmp_db) + for tbl in ("code", "document", "paper", "image", "concept", "rationale"): + rows = _query(conn, f"MATCH (n:{tbl}) RETURN count(n)") + assert rows == [[0]] + _close(db, conn) + + +# --- ingest_extraction: CREATE mode --- + +def test_ingest_extraction_create_mode(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + rows = _query(conn, "MATCH (n:code) RETURN n.id ORDER BY n.id") + ids = sorted([r[0] for r in rows]) + assert "n_attention" in ids + assert "n_transformer" in ids + assert "n_layernorm" in ids + edge_rows = _query(conn, "MATCH (a:code)-[e:edge_code_code]->(b:code) RETURN count(e)") + assert edge_rows[0][0] == 2 + _close(db, conn) + + +# --- ingest_extraction: MERGE mode --- + +def test_ingest_extraction_merge_mode(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + ext["nodes"][0]["label"] = "TransformerV2" + ingest_extraction(conn, ext, incremental=True) + rows = _query(conn, "MATCH (n:code) WHERE n.id = 'n_transformer' RETURN n.label") + assert rows[0][0] == "TransformerV2" + count = _query(conn, "MATCH (n:code) RETURN count(n)") + assert count[0][0] == 3 + _close(db, conn) + + +# --- file_type routing --- + +def test_ingest_extraction_file_type_routing(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + doc_rows = _query(conn, "MATCH (n:document) RETURN n.id") + assert len(doc_rows) == 1 + assert doc_rows[0][0] == "n_concept_attn" + _close(db, conn) + + +# --- prune_sources --- + +def test_ingest_extraction_prune(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + before = _query(conn, "MATCH (n:code) RETURN count(n)")[0][0] + assert before == 3 + ingest_extraction(conn, ext, incremental=True, prune_sources=["model.py"]) + after_prune = _query(conn, "MATCH (n:code) RETURN count(n)")[0][0] + assert after_prune == 3 + _close(db, conn) + + +# --- fallback rel table --- + +def test_fallback_rel_table(tmp_db): + from graphify.storage import ingest_extraction, _ensure_rel_table, _created_rel_tables + db, conn = _init(tmp_db) + assert ("paper", "document") not in _created_rel_tables + _ensure_rel_table(conn, "paper", "document") + assert ("paper", "document") in _created_rel_tables + _close(db, conn) + + +# --- communities --- + +def test_ingest_communities(tmp_db): + from graphify.storage import ingest_extraction, ingest_communities + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + communities = {0: ["n_transformer", "n_attention"], 1: ["n_layernorm"]} + ingest_communities(conn, communities) + rows = _query(conn, "MATCH (n:code) WHERE n.id = 'n_transformer' RETURN n.community") + assert rows[0][0] == 0 + rows = _query(conn, "MATCH (n:code) WHERE n.id = 'n_layernorm' RETURN n.community") + assert rows[0][0] == 1 + _close(db, conn) + + +# --- execute_cypher --- + +def test_execute_cypher(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + rows = _query(conn, "MATCH (n:code) RETURN n.label ORDER BY n.id") + labels = [r[0] for r in rows] + assert "MultiHeadAttention" in labels + assert "Transformer" in labels + _close(db, conn) + + +def test_execute_cypher_bad_query(tmp_db): + db, conn = _init(tmp_db) + with pytest.raises(RuntimeError): + _query(conn, "THIS IS NOT VALID CYPHER") + _close(db, conn) + + +# --- roundtrip consistency --- + +def test_roundtrip_node_count(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + total = 0 + for tbl in ("code", "document", "paper", "image", "concept", "rationale"): + rows = _query(conn, f"MATCH (n:{tbl}) RETURN count(n)") + total += rows[0][0] + assert total == len(ext["nodes"]) + _close(db, conn) From 2a8824593261b72180da419091ee814623fafd42 Mon Sep 17 00:00:00 2001 From: BingqingLyu Date: Tue, 26 May 2026 15:10:19 +0800 Subject: [PATCH 02/10] fix: split edge tables by relation to avoid NeuG 4096-row overflow Edge tables are now named edge_{src}_{tgt}_{relation} instead of edge_{src}_{tgt}. This keeps each table well under NeuG 0.1.0s 4096-row-per-table limit (max single table ~2475 rows for calls). Removes EXTRACTED/INFERRED routing distinction -- all edges are uniformly routed by their relation field. Co-Authored-By: Claude Opus 4.6 --- graphify/storage.py | 98 +++++++++++++++++++++--------------------- tests/test_neug_e2e.sh | 6 +-- tests/test_storage.py | 10 ++--- 3 files changed, 58 insertions(+), 56 deletions(-) diff --git a/graphify/storage.py b/graphify/storage.py index 3c15c3fc2e..f71342896b 100644 --- a/graphify/storage.py +++ b/graphify/storage.py @@ -42,32 +42,38 @@ } # --------------------------------------------------------------------------- -# Pre-built edge tables (based on actual graphify data model) +# Edge tables — split by (src_type, tgt_type, relation). +# Keeps each table under NeuG 0.1.0's 4096-row limit. # --------------------------------------------------------------------------- -_EDGE_PAIRS: list[tuple[str, str]] = [ - ("code", "code"), - ("code", "rationale"), - ("code", "concept"), - ("code", "document"), - ("document", "document"), - ("document", "concept"), - ("paper", "concept"), - ("paper", "paper"), - ("concept", "concept"), - ("image", "concept"), -] - -_EDGE_DDL_TEMPLATE = """CREATE REL TABLE IF NOT EXISTS edge_{src}_{tgt}( +_EDGE_DDL_TEMPLATE = """CREATE REL TABLE IF NOT EXISTS {tbl}( FROM {src} TO {tgt}, relation STRING, confidence STRING, confidence_score DOUBLE, source_file STRING, weight DOUBLE)""" -_created_rel_tables: set[tuple[str, str]] = set() +# Known relation types per (src, tgt) pair — pre-built at init time. +_KNOWN_RELATIONS: dict[tuple[str, str], list[str]] = { + ("code", "code"): [ + "calls", "contains", "method", "uses", "inherits", "defines", + "references", "imports", "imports_from", "listened_by", "case_of", + "references_constant", "bound_to", "uses_static_prop", "uses_config", + ], + ("rationale", "code"): ["rationale_for"], +} + +_created_rel_tables: set[str] = set() + +def _sanitize_rel_name(relation: str) -> str: + """Normalize a relation string into a safe table-name suffix.""" + r = relation.lower().strip() + r = re.sub(r"[^a-z0-9_]", "_", r) + r = re.sub(r"_+", "_", r).strip("_") + return r or "rel" -def _edge_table_name(src_type: str, tgt_type: str) -> str: - return f"edge_{src_type}_{tgt_type}" + +def _edge_table_name(src_type: str, tgt_type: str, relation: str) -> str: + return f"edge_{src_type}_{tgt_type}_{_sanitize_rel_name(relation)}" # --------------------------------------------------------------------------- @@ -100,20 +106,24 @@ def init_db(db_path: str) -> tuple[neug.Database, object]: conn.execute(ddl) _created_rel_tables.clear() - for src, tgt in _EDGE_PAIRS: - conn.execute(_EDGE_DDL_TEMPLATE.format(src=src, tgt=tgt)) - _created_rel_tables.add((src, tgt)) + + for (src, tgt), rels in _KNOWN_RELATIONS.items(): + for rel in rels: + tbl = _edge_table_name(src, tgt, rel) + conn.execute(_EDGE_DDL_TEMPLATE.format(tbl=tbl, src=src, tgt=tgt)) + _created_rel_tables.add(tbl) return db, conn -def _ensure_rel_table(conn: object, src_type: str, tgt_type: str) -> None: - """Create an edge table on-the-fly if not already pre-built.""" - pair = (src_type, tgt_type) - if pair in _created_rel_tables: - return - conn.execute(_EDGE_DDL_TEMPLATE.format(src=src_type, tgt=tgt_type)) - _created_rel_tables.add(pair) +def _ensure_rel_table(conn: object, src_type: str, tgt_type: str, relation: str) -> str: + """Resolve edge table name, creating on-the-fly if needed. Returns table name.""" + tbl = _edge_table_name(src_type, tgt_type, relation) + if tbl in _created_rel_tables: + return tbl + conn.execute(_EDGE_DDL_TEMPLATE.format(tbl=tbl, src=src_type, tgt=tgt_type)) + _created_rel_tables.add(tbl) + return tbl def _fix_file_type(ft: str | None) -> str: @@ -227,31 +237,23 @@ def ingest_extraction( if not src_ft or not tgt_ft: continue - _ensure_rel_table(conn, src_ft, tgt_ft) - tbl = _edge_table_name(src_ft, tgt_ft) - rel = _cesc(edge.get("relation", "")) - conf = _cesc(edge.get("confidence", "")) + rel_raw = edge.get("relation", "") + conf_raw = edge.get("confidence", "") + tbl = _ensure_rel_table(conn, src_ft, tgt_ft, rel_raw) + rel = _cesc(rel_raw) + conf = _cesc(conf_raw) conf_score = float(edge.get("confidence_score", 0.0)) e_sf = _cesc(_norm_source_file(edge.get("source_file"), _root) or "") weight = float(edge.get("weight", 1.0)) try: - if incremental: - conn.execute( - f"MATCH (a:{src_ft} {{id: '{_cesc(src_id)}'}}), " - f"(b:{tgt_ft} {{id: '{_cesc(tgt_id)}'}}) " - f"CREATE (a)-[:{tbl} {{relation: '{rel}', confidence: '{conf}', " - f"confidence_score: {conf_score}, source_file: '{e_sf}', " - f"weight: {weight}}}]->(b)" - ) - else: - conn.execute( - f"MATCH (a:{src_ft} {{id: '{_cesc(src_id)}'}}), " - f"(b:{tgt_ft} {{id: '{_cesc(tgt_id)}'}}) " - f"CREATE (a)-[:{tbl} {{relation: '{rel}', confidence: '{conf}', " - f"confidence_score: {conf_score}, source_file: '{e_sf}', " - f"weight: {weight}}}]->(b)" - ) + conn.execute( + f"MATCH (a:{src_ft} {{id: '{_cesc(src_id)}'}}), " + f"(b:{tgt_ft} {{id: '{_cesc(tgt_id)}'}}) " + f"CREATE (a)-[:{tbl} {{relation: '{rel}', confidence: '{conf}', " + f"confidence_score: {conf_score}, source_file: '{e_sf}', " + f"weight: {weight}}}]->(b)" + ) except RuntimeError: _e_errors += 1 diff --git a/tests/test_neug_e2e.sh b/tests/test_neug_e2e.sh index 48b5a246f8..983b94fb20 100755 --- a/tests/test_neug_e2e.sh +++ b/tests/test_neug_e2e.sh @@ -204,7 +204,7 @@ else fi # Query edges - CYPHER_EDGES=$(python3 -m graphify cypher "MATCH (a:code)-[e:edge_code_code]->(b:code) RETURN a.label, e.relation, b.label LIMIT 5" --db "$DB_PATH" 2>/dev/null | grep -v "^INFO\|^E20") + CYPHER_EDGES=$(python3 -m graphify cypher "MATCH (a:code)-[e]->(b:code) RETURN a.label, e.relation, b.label LIMIT 5" --db "$DB_PATH" 2>/dev/null | grep -v "^INFO\|^E20") if [ -n "$CYPHER_EDGES" ]; then pass "cypher edge query returned results" else @@ -298,8 +298,8 @@ ingest_extraction(conn, extraction, incremental=False) nodes = execute_cypher(conn, "MATCH (n:code) RETURN count(n)") assert nodes[0][0] == 2, f"Expected 2 code nodes, got {nodes[0][0]}" -edges = execute_cypher(conn, "MATCH ()-[e:edge_code_code]->() RETURN count(e)") -assert edges[0][0] == 1, f"Expected 1 code-code edge, got {edges[0][0]}" +edges = execute_cypher(conn, "MATCH ()-[e:edge_code_code_calls]->() RETURN count(e)") +assert edges[0][0] == 1, f"Expected 1 code-code-calls edge, got {edges[0][0]}" # Communities ingest_communities(conn, {0: ["fn_main", "fn_helper"], 1: ["concept_arch"]}) diff --git a/tests/test_storage.py b/tests/test_storage.py index 839d9c9abd..b9024ac639 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -65,7 +65,7 @@ def test_ingest_extraction_create_mode(tmp_db): assert "n_attention" in ids assert "n_transformer" in ids assert "n_layernorm" in ids - edge_rows = _query(conn, "MATCH (a:code)-[e:edge_code_code]->(b:code) RETURN count(e)") + edge_rows = _query(conn, "MATCH (a:code)-[e:edge_code_code_contains]->(b:code) RETURN count(e)") assert edge_rows[0][0] == 2 _close(db, conn) @@ -117,11 +117,11 @@ def test_ingest_extraction_prune(tmp_db): # --- fallback rel table --- def test_fallback_rel_table(tmp_db): - from graphify.storage import ingest_extraction, _ensure_rel_table, _created_rel_tables + from graphify.storage import _ensure_rel_table, _created_rel_tables db, conn = _init(tmp_db) - assert ("paper", "document") not in _created_rel_tables - _ensure_rel_table(conn, "paper", "document") - assert ("paper", "document") in _created_rel_tables + tbl = _ensure_rel_table(conn, "paper", "document", "cites") + assert tbl == "edge_paper_document_cites" + assert tbl in _created_rel_tables _close(db, conn) From 6d233004d549dd7e617639cefb98c5cff5f29dd7 Mon Sep 17 00:00:00 2001 From: BingqingLyu Date: Wed, 27 May 2026 15:53:43 +0800 Subject: [PATCH 03/10] fix: use cache_root for id_remap so incremental extraction produces stable file node IDs When extracting a single file incrementally, the auto-inferred root (paths[0].parent) differs from the project root used during full extraction, causing file node IDs to mismatch (e.g. "build_py" vs "graphify_build_py"). This created duplicate file nodes on each incremental update (+2 instead of +1). Fix: use cache_root (the project target directory passed by __main__.py) for relative_to() in the id_remap step, ensuring file node IDs are consistent regardless of whether extraction is full or incremental. Co-Authored-By: Claude Opus 4.6 --- graphify/extract.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/graphify/extract.py b/graphify/extract.py index c2443f1b12..a4a25fbb3f 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -9087,11 +9087,12 @@ def extract( # Remap file node IDs from absolute-path-derived to project-relative so # graph.json edge endpoints are stable across machines (#502) + remap_root = (cache_root or root).resolve() if cache_root else root id_remap: dict[str, str] = {} for path in paths: old_id = _make_id(str(path)) try: - new_id = _make_id(str(path.relative_to(root))) + new_id = _make_id(str(path.resolve().relative_to(remap_root))) except ValueError: continue if old_id != new_id: From 96ee7df29ce7850eab921133d4fea388d977559d Mon Sep 17 00:00:00 2001 From: BingqingLyu Date: Thu, 28 May 2026 10:56:58 +0800 Subject: [PATCH 04/10] refactor: split init_db into open-only + ensure_schema for schema creation Separates database opening (init_db) from schema DDL (ensure_schema) so read-only consumers (CLI cypher, MCP server) can open an existing graph.db without re-running CREATE TABLE statements. Co-Authored-By: Claude Opus 4.6 --- graphify/__main__.py | 6 ++++-- graphify/storage.py | 9 +++++---- tests/test_cypher_cli.py | 3 ++- tests/test_storage.py | 6 ++++-- 4 files changed, 15 insertions(+), 9 deletions(-) diff --git a/graphify/__main__.py b/graphify/__main__.py index 2d6b769027..9f9ec5162a 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -3342,10 +3342,11 @@ def _progress(idx: int, total: int, _result: dict) -> None: json.dumps(merged, indent=2), encoding="utf-8" ) try: - from graphify.storage import init_db as _init_db, ingest_extraction as _ingest, close_db as _close_db + from graphify.storage import init_db as _init_db, ensure_schema as _ensure_schema, ingest_extraction as _ingest, close_db as _close_db _db_path = str(graphify_out / "graph.db") _is_inc = Path(_db_path).exists() _db, _conn = _init_db(_db_path) + _ensure_schema(_conn) _ingest(_conn, merged, incremental=_is_inc, prune_sources=deleted_files or None, root=target) _close_db(_db, _conn) @@ -3432,10 +3433,11 @@ def _progress(idx: int, total: int, _result: dict) -> None: _backup(graphify_out) _to_json(G, communities, str(graph_json_path), force=True) try: - from graphify.storage import init_db as _init_db, ingest_extraction as _ingest, ingest_communities as _ingest_comm, close_db as _close_db + from graphify.storage import init_db as _init_db, ensure_schema as _ensure_schema, ingest_extraction as _ingest, ingest_communities as _ingest_comm, close_db as _close_db _db_path = str(graphify_out / "graph.db") _is_inc = Path(_db_path).exists() _db, _conn = _init_db(_db_path) + _ensure_schema(_conn) _ingest(_conn, merged, incremental=_is_inc, prune_sources=deleted_files or None, root=target) _ingest_comm(_conn, communities) diff --git a/graphify/storage.py b/graphify/storage.py index f71342896b..f78e05fc44 100644 --- a/graphify/storage.py +++ b/graphify/storage.py @@ -43,7 +43,6 @@ # --------------------------------------------------------------------------- # Edge tables — split by (src_type, tgt_type, relation). -# Keeps each table under NeuG 0.1.0's 4096-row limit. # --------------------------------------------------------------------------- _EDGE_DDL_TEMPLATE = """CREATE REL TABLE IF NOT EXISTS {tbl}( @@ -98,10 +97,14 @@ def _cesc(value: str) -> str: def init_db(db_path: str) -> tuple[neug.Database, object]: - """Create or open a NeuG database and initialize schema.""" + """Open (or create) a NeuG database and connect.""" db = neug.Database(db_path) conn = db.connect() + return db, conn + +def ensure_schema(conn: object) -> None: + """Create node/edge tables if not exist. Call once during extract.""" for ddl in _NODE_TABLES.values(): conn.execute(ddl) @@ -113,8 +116,6 @@ def init_db(db_path: str) -> tuple[neug.Database, object]: conn.execute(_EDGE_DDL_TEMPLATE.format(tbl=tbl, src=src, tgt=tgt)) _created_rel_tables.add(tbl) - return db, conn - def _ensure_rel_table(conn: object, src_type: str, tgt_type: str, relation: str) -> str: """Resolve edge table name, creating on-the-fly if needed. Returns table name.""" diff --git a/tests/test_cypher_cli.py b/tests/test_cypher_cli.py index fb705d9b5d..d8c1065eaf 100644 --- a/tests/test_cypher_cli.py +++ b/tests/test_cypher_cli.py @@ -20,10 +20,11 @@ def _build_db(tmp_path) -> str: - from graphify.storage import init_db, ingest_extraction, close_db + from graphify.storage import init_db, ensure_schema, ingest_extraction, close_db db_path = str(tmp_path / "graph.db") ext = json.loads(EXTRACTION_JSON.read_text()) db, conn = init_db(db_path) + ensure_schema(conn) ingest_extraction(conn, ext, incremental=False) close_db(db, conn) return db_path diff --git a/tests/test_storage.py b/tests/test_storage.py index b9024ac639..c8b21c43dd 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -29,8 +29,10 @@ def tmp_db(tmp_path): def _init(db_path): - from graphify.storage import init_db - return init_db(db_path) + from graphify.storage import init_db, ensure_schema + db, conn = init_db(db_path) + ensure_schema(conn) + return db, conn def _close(db, conn): From 2dd3a1676bec5a2768c9f5fdd306cefbcecb1b22 Mon Sep 17 00:00:00 2001 From: BingqingLyu Date: Thu, 28 May 2026 15:46:01 +0800 Subject: [PATCH 05/10] feat: use native MERGE for incremental updates, skip DDL on existing db - Replace MATCH+check+SET workaround with standard Cypher MERGE ON CREATE/ON MATCH syntax (now supported in NeuG 0.1.2) - ensure_schema(create_tables=False) skips DDL on incremental runs, avoiding "table already exists" noise from the C++ layer Co-Authored-By: Claude Opus 4.6 --- graphify/__main__.py | 4 ++-- graphify/storage.py | 53 ++++++++++++++++++++++---------------------- 2 files changed, 29 insertions(+), 28 deletions(-) diff --git a/graphify/__main__.py b/graphify/__main__.py index 9f9ec5162a..dfa4c1a3e7 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -3346,7 +3346,7 @@ def _progress(idx: int, total: int, _result: dict) -> None: _db_path = str(graphify_out / "graph.db") _is_inc = Path(_db_path).exists() _db, _conn = _init_db(_db_path) - _ensure_schema(_conn) + _ensure_schema(_conn, create_tables=not _is_inc) _ingest(_conn, merged, incremental=_is_inc, prune_sources=deleted_files or None, root=target) _close_db(_db, _conn) @@ -3437,7 +3437,7 @@ def _progress(idx: int, total: int, _result: dict) -> None: _db_path = str(graphify_out / "graph.db") _is_inc = Path(_db_path).exists() _db, _conn = _init_db(_db_path) - _ensure_schema(_conn) + _ensure_schema(_conn, create_tables=not _is_inc) _ingest(_conn, merged, incremental=_is_inc, prune_sources=deleted_files or None, root=target) _ingest_comm(_conn, communities) diff --git a/graphify/storage.py b/graphify/storage.py index f78e05fc44..3faa873c47 100644 --- a/graphify/storage.py +++ b/graphify/storage.py @@ -103,17 +103,24 @@ def init_db(db_path: str) -> tuple[neug.Database, object]: return db, conn -def ensure_schema(conn: object) -> None: - """Create node/edge tables if not exist. Call once during extract.""" - for ddl in _NODE_TABLES.values(): - conn.execute(ddl) +def ensure_schema(conn: object, *, create_tables: bool = True) -> None: + """Populate known table registry; optionally execute DDL. + create_tables=True (first build): run CREATE TABLE statements. + create_tables=False (incremental): only populate _created_rel_tables set + so _ensure_rel_table() knows what exists. + """ _created_rel_tables.clear() + if create_tables: + for ddl in _NODE_TABLES.values(): + conn.execute(ddl) + for (src, tgt), rels in _KNOWN_RELATIONS.items(): for rel in rels: tbl = _edge_table_name(src, tgt, rel) - conn.execute(_EDGE_DDL_TEMPLATE.format(tbl=tbl, src=src, tgt=tgt)) + if create_tables: + conn.execute(_EDGE_DDL_TEMPLATE.format(tbl=tbl, src=src, tgt=tgt)) _created_rel_tables.add(tbl) @@ -182,30 +189,24 @@ def ingest_extraction( try: if incremental: - existing = list(conn.execute( - f"MATCH (n:{ft} {{id: '{_cesc(nid)}'}}) RETURN n.id" - )) - if existing: - props = f"n.label = '{_cesc(label)}', n.source_file = '{_cesc(sf)}'" - if ft == "code": - props += f", n.source_location = '{_cesc(sl)}'" + if ft == "code": conn.execute( - f"MATCH (n:{ft} {{id: '{_cesc(nid)}'}}) SET {props}" + f"MERGE (n:code {{id: '{_cesc(nid)}'}}) " + f"ON CREATE SET n.label = '{_cesc(label)}', " + f"n.source_file = '{_cesc(sf)}', " + f"n.source_location = '{_cesc(sl)}' " + f"ON MATCH SET n.label = '{_cesc(label)}', " + f"n.source_file = '{_cesc(sf)}', " + f"n.source_location = '{_cesc(sl)}'" ) else: - if ft == "code": - conn.execute( - f"CREATE (n:code {{id: '{_cesc(nid)}', " - f"label: '{_cesc(label)}', " - f"source_file: '{_cesc(sf)}', " - f"source_location: '{_cesc(sl)}'}})" - ) - else: - conn.execute( - f"CREATE (n:{ft} {{id: '{_cesc(nid)}', " - f"label: '{_cesc(label)}', " - f"source_file: '{_cesc(sf)}'}})" - ) + conn.execute( + f"MERGE (n:{ft} {{id: '{_cesc(nid)}'}}) " + f"ON CREATE SET n.label = '{_cesc(label)}', " + f"n.source_file = '{_cesc(sf)}' " + f"ON MATCH SET n.label = '{_cesc(label)}', " + f"n.source_file = '{_cesc(sf)}'" + ) else: if ft == "code": conn.execute( From f18fa2cbf92354193191f7cfd499b75881e75c96 Mon Sep 17 00:00:00 2001 From: BingqingLyu Date: Thu, 28 May 2026 16:06:03 +0800 Subject: [PATCH 06/10] chore: require neug>=0.1.2 for MERGE support Co-Authored-By: Claude Opus 4.6 --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b160f8ddf0..ad1390968f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,8 +65,8 @@ gemini = ["openai", "tiktoken"] openai = ["openai", "tiktoken"] chinese = ["jieba"] sql = ["tree-sitter-sql"] -neug = ["neug"] -all = ["mcp", "neo4j", "neug", "pypdf", "markdownify", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper", "yt-dlp", "matplotlib", "openai", "tiktoken", "boto3", "tree-sitter-sql", "jieba"] +neug = ["neug>=0.1.2"] +all = ["mcp", "neo4j", "neug>=0.1.2", "pypdf", "markdownify", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper", "yt-dlp", "matplotlib", "openai", "tiktoken", "boto3", "tree-sitter-sql", "jieba"] [project.scripts] graphify = "graphify.__main__:main" From 5fa77f9ebbd9aea6b31e82c21fa39011d71bbdc4 Mon Sep 17 00:00:00 2001 From: BingqingLyu Date: Fri, 29 May 2026 10:49:01 +0800 Subject: [PATCH 07/10] address PR #1056 review: parameterised queries, lazy import, per-conn registry - Replace all _cesc() string interpolation with NeuG native $param syntax to prevent Cypher injection (community SET uses int literal due to NeuG limitation on parameterised SET) - Move `import neug` into init_db() for lazy loading - Make rel table registry per-connection via ensure_schema() return value - ingest_extraction() returns node_types dict for O(n) community writes - Require neug>=0.1.2,<0.2 for MERGE support - Remove tests/test_neug_e2e.sh (manual script, not automated) Co-Authored-By: Claude Opus 4.6 --- graphify/__main__.py | 14 +- graphify/storage.py | 159 ++++++++------- pyproject.toml | 4 +- tests/test_cypher_cli.py | 4 +- tests/test_neug_e2e.sh | 418 --------------------------------------- tests/test_storage.py | 7 +- 6 files changed, 106 insertions(+), 500 deletions(-) delete mode 100755 tests/test_neug_e2e.sh diff --git a/graphify/__main__.py b/graphify/__main__.py index dfa4c1a3e7..a77c0e993b 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -3346,9 +3346,10 @@ def _progress(idx: int, total: int, _result: dict) -> None: _db_path = str(graphify_out / "graph.db") _is_inc = Path(_db_path).exists() _db, _conn = _init_db(_db_path) - _ensure_schema(_conn, create_tables=not _is_inc) + _known = _ensure_schema(_conn, create_tables=not _is_inc) _ingest(_conn, merged, incremental=_is_inc, - prune_sources=deleted_files or None, root=target) + prune_sources=deleted_files or None, root=target, + known_tables=_known) _close_db(_db, _conn) print("[graphify extract] graph.db written (powered by NeuG)") except ImportError: @@ -3437,10 +3438,11 @@ def _progress(idx: int, total: int, _result: dict) -> None: _db_path = str(graphify_out / "graph.db") _is_inc = Path(_db_path).exists() _db, _conn = _init_db(_db_path) - _ensure_schema(_conn, create_tables=not _is_inc) - _ingest(_conn, merged, incremental=_is_inc, - prune_sources=deleted_files or None, root=target) - _ingest_comm(_conn, communities) + _known = _ensure_schema(_conn, create_tables=not _is_inc) + _ntypes = _ingest(_conn, merged, incremental=_is_inc, + prune_sources=deleted_files or None, root=target, + known_tables=_known) + _ingest_comm(_conn, communities, node_types=_ntypes) _close_db(_db, _conn) print("[graphify extract] graph.db written (powered by NeuG)") except ImportError: diff --git a/graphify/storage.py b/graphify/storage.py index 3faa873c47..ce55a8539e 100644 --- a/graphify/storage.py +++ b/graphify/storage.py @@ -1,18 +1,19 @@ """NeuG graph database adapter for graphify. Provides an optional parallel storage engine alongside NetworkX. -All functions are guarded by `import neug` — when NeuG is not installed, -callers should catch ImportError and skip silently. +NeuG is lazily imported — when not installed, callers should catch +ImportError at the call site and skip silently. + +All property values interpolated into Cypher statements use NeuG's native +parameterised queries ($param syntax) to prevent injection. Table/label +names (which come from a fixed internal set, not user input) are still +interpolated as identifiers. """ from __future__ import annotations -import os import re -import unicodedata from pathlib import Path -import neug - from .build import _FILE_TYPE_SYNONYMS, _normalize_id, _norm_source_file from .validate import VALID_FILE_TYPES @@ -60,8 +61,6 @@ ("rationale", "code"): ["rationale_for"], } -_created_rel_tables: set[str] = set() - def _sanitize_rel_name(relation: str) -> str: """Normalize a relation string into a safe table-name suffix.""" @@ -75,42 +74,32 @@ def _edge_table_name(src_type: str, tgt_type: str, relation: str) -> str: return f"edge_{src_type}_{tgt_type}_{_sanitize_rel_name(relation)}" -# --------------------------------------------------------------------------- -# Cypher string escaping -# --------------------------------------------------------------------------- - - -def _cesc(value: str) -> str: - return ( - value - .replace("\\", "\\\\") - .replace("'", "\\'") - .replace("\n", "\\n") - .replace("\r", "\\r") - .replace("\t", "\\t") - ) - - # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- -def init_db(db_path: str) -> tuple[neug.Database, object]: - """Open (or create) a NeuG database and connect.""" +def init_db(db_path: str) -> tuple: + """Open (or create) a NeuG database and connect. + + Returns (db, conn). Raises ImportError if neug is not installed. + """ + import neug db = neug.Database(db_path) conn = db.connect() return db, conn -def ensure_schema(conn: object, *, create_tables: bool = True) -> None: +def ensure_schema(conn: object, *, create_tables: bool = True) -> set[str]: """Populate known table registry; optionally execute DDL. create_tables=True (first build): run CREATE TABLE statements. - create_tables=False (incremental): only populate _created_rel_tables set + create_tables=False (incremental): only build the registry set so _ensure_rel_table() knows what exists. + + Returns the set of known rel table names (per-connection registry). """ - _created_rel_tables.clear() + created: set[str] = set() if create_tables: for ddl in _NODE_TABLES.values(): @@ -121,16 +110,21 @@ def ensure_schema(conn: object, *, create_tables: bool = True) -> None: tbl = _edge_table_name(src, tgt, rel) if create_tables: conn.execute(_EDGE_DDL_TEMPLATE.format(tbl=tbl, src=src, tgt=tgt)) - _created_rel_tables.add(tbl) + created.add(tbl) + return created -def _ensure_rel_table(conn: object, src_type: str, tgt_type: str, relation: str) -> str: + +def _ensure_rel_table( + conn: object, src_type: str, tgt_type: str, relation: str, + known: set[str], +) -> str: """Resolve edge table name, creating on-the-fly if needed. Returns table name.""" tbl = _edge_table_name(src_type, tgt_type, relation) - if tbl in _created_rel_tables: + if tbl in known: return tbl conn.execute(_EDGE_DDL_TEMPLATE.format(tbl=tbl, src=src_type, tgt=tgt_type)) - _created_rel_tables.add(tbl) + known.add(tbl) return tbl @@ -148,13 +142,17 @@ def ingest_extraction( incremental: bool = False, prune_sources: list[str] | None = None, root: str | Path | None = None, -) -> None: + known_tables: set[str] | None = None, +) -> dict[str, str]: """Write an extraction dict into NeuG. incremental=False: first build — uses CREATE (faster). incremental=True: update — uses MERGE (upsert). + + Returns node_types dict (id -> file_type) for use by ingest_communities. """ _root = str(Path(root).resolve()) if root else None + _known = known_tables if known_tables is not None else set() # --- prune deleted/changed files first --- if prune_sources: @@ -162,8 +160,8 @@ def ingest_extraction( sf_norm = _norm_source_file(sf, _root) or sf for tbl in _NODE_TABLES: conn.execute( - f"MATCH (n:{tbl}) WHERE n.source_file = '{_cesc(sf_norm)}' " - f"DETACH DELETE n" + f"MATCH (n:{tbl}) WHERE n.source_file = $sf DETACH DELETE n", + parameters={"sf": sf_norm}, ) # --- build node lookup: id -> file_type --- @@ -191,35 +189,32 @@ def ingest_extraction( if incremental: if ft == "code": conn.execute( - f"MERGE (n:code {{id: '{_cesc(nid)}'}}) " - f"ON CREATE SET n.label = '{_cesc(label)}', " - f"n.source_file = '{_cesc(sf)}', " - f"n.source_location = '{_cesc(sl)}' " - f"ON MATCH SET n.label = '{_cesc(label)}', " - f"n.source_file = '{_cesc(sf)}', " - f"n.source_location = '{_cesc(sl)}'" + f"MERGE (n:code {{id: $nid}}) " + f"ON CREATE SET n.label = $label, " + f"n.source_file = $sf, n.source_location = $sl " + f"ON MATCH SET n.label = $label, " + f"n.source_file = $sf, n.source_location = $sl", + parameters={"nid": nid, "label": label, "sf": sf, "sl": sl}, ) else: conn.execute( - f"MERGE (n:{ft} {{id: '{_cesc(nid)}'}}) " - f"ON CREATE SET n.label = '{_cesc(label)}', " - f"n.source_file = '{_cesc(sf)}' " - f"ON MATCH SET n.label = '{_cesc(label)}', " - f"n.source_file = '{_cesc(sf)}'" + f"MERGE (n:{ft} {{id: $nid}}) " + f"ON CREATE SET n.label = $label, n.source_file = $sf " + f"ON MATCH SET n.label = $label, n.source_file = $sf", + parameters={"nid": nid, "label": label, "sf": sf}, ) else: if ft == "code": conn.execute( - f"CREATE (n:code {{id: '{_cesc(nid)}', " - f"label: '{_cesc(label)}', " - f"source_file: '{_cesc(sf)}', " - f"source_location: '{_cesc(sl)}'}})" + f"CREATE (n:code {{id: $nid, label: $label, " + f"source_file: $sf, source_location: $sl}})", + parameters={"nid": nid, "label": label, "sf": sf, "sl": sl}, ) else: conn.execute( - f"CREATE (n:{ft} {{id: '{_cesc(nid)}', " - f"label: '{_cesc(label)}', " - f"source_file: '{_cesc(sf)}'}})" + f"CREATE (n:{ft} {{id: $nid, label: $label, " + f"source_file: $sf}})", + parameters={"nid": nid, "label": label, "sf": sf}, ) except RuntimeError: _n_errors += 1 @@ -241,20 +236,24 @@ def ingest_extraction( rel_raw = edge.get("relation", "") conf_raw = edge.get("confidence", "") - tbl = _ensure_rel_table(conn, src_ft, tgt_ft, rel_raw) - rel = _cesc(rel_raw) - conf = _cesc(conf_raw) + tbl = _ensure_rel_table(conn, src_ft, tgt_ft, rel_raw, _known) conf_score = float(edge.get("confidence_score", 0.0)) - e_sf = _cesc(_norm_source_file(edge.get("source_file"), _root) or "") + e_sf = _norm_source_file(edge.get("source_file"), _root) or "" weight = float(edge.get("weight", 1.0)) try: conn.execute( - f"MATCH (a:{src_ft} {{id: '{_cesc(src_id)}'}}), " - f"(b:{tgt_ft} {{id: '{_cesc(tgt_id)}'}}) " - f"CREATE (a)-[:{tbl} {{relation: '{rel}', confidence: '{conf}', " - f"confidence_score: {conf_score}, source_file: '{e_sf}', " - f"weight: {weight}}}]->(b)" + f"MATCH (a:{src_ft} {{id: $src_id}}), " + f"(b:{tgt_ft} {{id: $tgt_id}}) " + f"CREATE (a)-[:{tbl} {{relation: $rel, confidence: $conf, " + f"confidence_score: $conf_score, source_file: $e_sf, " + f"weight: $weight}}]->(b)", + parameters={ + "src_id": src_id, "tgt_id": tgt_id, + "rel": rel_raw, "conf": conf_raw, + "conf_score": conf_score, "e_sf": e_sf, + "weight": weight, + }, ) except RuntimeError: _e_errors += 1 @@ -266,23 +265,45 @@ def ingest_extraction( _n_errors, _e_errors, ) + return node_types + def ingest_communities( conn: object, communities: dict[int, list[str]], community_labels: dict[int, str] | None = None, + node_types: dict[str, str] | None = None, ) -> None: - """Write community assignments into NeuG node properties.""" + """Write community assignments into NeuG node properties. + + If node_types is provided (id -> file_type mapping from ingest_extraction), + each node is looked up in its specific table directly. Otherwise falls + back to probing all 6 tables (slower). + + Note: NeuG does not support parameterised SET for non-string values, + so community ID is interpolated as an integer literal. The id value + uses a parameterised query. + """ for cid, node_ids in communities.items(): + cid_int = int(cid) for nid in node_ids: nid_norm = _normalize_id(nid) if not nid_norm: continue - for tbl in _NODE_TABLES: + if node_types and nid_norm in node_types: + tbl = node_types[nid_norm] conn.execute( - f"MATCH (n:{tbl}) WHERE n.id = '{_cesc(nid_norm)}' " - f"SET n.community = {int(cid)}" + f"MATCH (n:{tbl}) WHERE n.id = $nid " + f"SET n.community = {cid_int}", + parameters={"nid": nid_norm}, ) + else: + for tbl in _NODE_TABLES: + conn.execute( + f"MATCH (n:{tbl}) WHERE n.id = $nid " + f"SET n.community = {cid_int}", + parameters={"nid": nid_norm}, + ) def execute_cypher(conn: object, query: str) -> list[list]: @@ -293,7 +314,7 @@ def execute_cypher(conn: object, query: str) -> list[list]: raise RuntimeError(f"Cypher query failed: {exc}") from exc -def close_db(db: neug.Database, conn: object) -> None: +def close_db(db: object, conn: object) -> None: """Close the NeuG connection and database.""" conn.close() db.close() diff --git a/pyproject.toml b/pyproject.toml index ad1390968f..f952bc8473 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,8 +65,8 @@ gemini = ["openai", "tiktoken"] openai = ["openai", "tiktoken"] chinese = ["jieba"] sql = ["tree-sitter-sql"] -neug = ["neug>=0.1.2"] -all = ["mcp", "neo4j", "neug>=0.1.2", "pypdf", "markdownify", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper", "yt-dlp", "matplotlib", "openai", "tiktoken", "boto3", "tree-sitter-sql", "jieba"] +neug = ["neug>=0.1.2,<0.2"] +all = ["mcp", "neo4j", "neug>=0.1.2,<0.2", "pypdf", "markdownify", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper", "yt-dlp", "matplotlib", "openai", "tiktoken", "boto3", "tree-sitter-sql", "jieba"] [project.scripts] graphify = "graphify.__main__:main" diff --git a/tests/test_cypher_cli.py b/tests/test_cypher_cli.py index d8c1065eaf..8e18a31250 100644 --- a/tests/test_cypher_cli.py +++ b/tests/test_cypher_cli.py @@ -24,8 +24,8 @@ def _build_db(tmp_path) -> str: db_path = str(tmp_path / "graph.db") ext = json.loads(EXTRACTION_JSON.read_text()) db, conn = init_db(db_path) - ensure_schema(conn) - ingest_extraction(conn, ext, incremental=False) + known = ensure_schema(conn) + ingest_extraction(conn, ext, incremental=False, known_tables=known) close_db(db, conn) return db_path diff --git a/tests/test_neug_e2e.sh b/tests/test_neug_e2e.sh deleted file mode 100755 index 983b94fb20..0000000000 --- a/tests/test_neug_e2e.sh +++ /dev/null @@ -1,418 +0,0 @@ -#!/usr/bin/env bash -# -# End-to-end test for graphify + NeuG integration. -# -# Tests the full flow: -# 1. graphify extract (AST-only, no LLM) → graph.json + graph.db -# 2. graphify cypher → query against graph.db -# 3. Incremental re-extract → graph.db updated -# 4. MCP server tool registration (smoke test) -# -# Prerequisites: -# - pip install neug -# - graphify installed from current source (pip install -e .) -# -# Usage: -# bash tests/test_neug_e2e.sh -# -set -uo pipefail - -RED='\033[0;31m' -GREEN='\033[0;32m' -YELLOW='\033[1;33m' -NC='\033[0m' - -PASS=0 -FAIL=0 -SKIP=0 - -pass() { echo -e " ${GREEN}PASS${NC}: $1"; ((PASS++)); } -fail() { echo -e " ${RED}FAIL${NC}: $1"; ((FAIL++)); } -skip() { echo -e " ${YELLOW}SKIP${NC}: $1"; ((SKIP++)); } - -echo "======================================" -echo " graphify + NeuG E2E Integration Test" -echo "======================================" -echo "" - -# --- Check prerequisites --- -if ! python3 -c "import neug" 2>/dev/null; then - echo "ERROR: neug not installed. Run: pip install neug" - exit 1 -fi - -if ! python3 -c "import graphify" 2>/dev/null; then - echo "ERROR: graphify not importable. Run: pip install -e . from graphify root" - exit 1 -fi - -# --- Setup test project --- -TMPDIR=$(mktemp -d) -trap "rm -rf $TMPDIR" EXIT - -PROJECT="$TMPDIR/sample_project" -mkdir -p "$PROJECT/src" - -cat > "$PROJECT/src/main.py" << 'PYEOF' -"""Main application module.""" - -class Database: - """Database connection manager.""" - def __init__(self, url: str): - self.url = url - self.conn = None - - def connect(self): - """Establish connection.""" - from src.utils import validate_url - validate_url(self.url) - self.conn = True - return self - - def query(self, sql: str): - """Execute a query.""" - if not self.conn: - raise RuntimeError("Not connected") - return [] - - -class App: - """Main application.""" - def __init__(self): - self.db = Database("localhost:5432") - - def run(self): - self.db.connect() - results = self.db.query("SELECT 1") - return results -PYEOF - -cat > "$PROJECT/src/utils.py" << 'PYEOF' -"""Utility functions.""" - -def validate_url(url: str) -> bool: - """Validate a database URL.""" - if not url: - raise ValueError("Empty URL") - return ":" in url - - -def format_result(row: dict) -> str: - """Format a query result row.""" - return ", ".join(f"{k}={v}" for k, v in row.items()) - - -class Logger: - """Simple logger.""" - def __init__(self, name: str): - self.name = name - - def info(self, msg: str): - print(f"[{self.name}] {msg}") -PYEOF - -cat > "$PROJECT/src/models.py" << 'PYEOF' -"""Data models.""" -from src.utils import Logger - -class User: - """User model.""" - def __init__(self, name: str, email: str): - self.name = name - self.email = email - self.logger = Logger("User") - - def save(self): - self.logger.info(f"Saving user {self.name}") - -class Session: - """Session model.""" - def __init__(self, user: User): - self.user = user - self.active = True - - def close(self): - self.active = False -PYEOF - -echo "Test project created at $PROJECT" -echo "" - -# ============================================================ -# TEST 1: First extract (AST-only, no cluster) -# ============================================================ -echo "[Test 1] graphify extract (first build, no-semantic, no-cluster)" - -OUTPUT=$(cd "$PROJECT" && GEMINI_API_KEY=dummy python3 -m graphify extract . --no-semantic --no-cluster 2>&1) - -CLEAN_OUTPUT=$(echo "$OUTPUT" | grep -v "^INFO\|^E20") -if echo "$CLEAN_OUTPUT" | grep -q "graph.db written"; then - pass "graph.db written message present" -else - if [ -d "$PROJECT/graphify-out/graph.db" ]; then - pass "graph.db created (message may be suppressed)" - else - fail "graph.db NOT created" - echo " Output: $(echo "$CLEAN_OUTPUT" | grep -i 'neug\|graph.db\|warning\|error' | head -5)" - fi -fi - -if [ -f "$PROJECT/graphify-out/graph.json" ]; then - pass "graph.json created" -else - fail "graph.json NOT created" -fi - -# Verify graph.json has nodes -NODE_COUNT=$(python3 -c " -import json -d = json.load(open('$PROJECT/graphify-out/graph.json')) -print(len(d.get('nodes', []))) -") -if [ "$NODE_COUNT" -gt 0 ]; then - pass "graph.json has $NODE_COUNT nodes" -else - fail "graph.json has 0 nodes" -fi - -echo "" - -# ============================================================ -# TEST 2: graphify cypher — query the database -# ============================================================ -echo "[Test 2] graphify cypher — Cypher queries against graph.db" - -DB_PATH="$PROJECT/graphify-out/graph.db" - -if [ ! -e "$DB_PATH" ]; then - skip "graph.db not available, skipping cypher tests" -else - # Count all nodes - CYPHER_COUNT=$(python3 -m graphify cypher "MATCH (n:code) RETURN count(n)" --db "$DB_PATH" 2>/dev/null | grep -v "^INFO\|^E20") - if [ -n "$CYPHER_COUNT" ] && [ "$CYPHER_COUNT" != "0" ]; then - pass "cypher count query returned: $CYPHER_COUNT" - else - fail "cypher count query returned empty or zero" - fi - - # Query node labels - CYPHER_LABELS=$(python3 -m graphify cypher "MATCH (n:code) RETURN n.label LIMIT 5" --db "$DB_PATH" 2>/dev/null | grep -v "^INFO\|^E20") - if [ -n "$CYPHER_LABELS" ]; then - pass "cypher label query returned results" - else - fail "cypher label query returned empty" - fi - - # Query edges - CYPHER_EDGES=$(python3 -m graphify cypher "MATCH (a:code)-[e]->(b:code) RETURN a.label, e.relation, b.label LIMIT 5" --db "$DB_PATH" 2>/dev/null | grep -v "^INFO\|^E20") - if [ -n "$CYPHER_EDGES" ]; then - pass "cypher edge query returned results" - else - fail "cypher edge query returned empty (may have no code->code edges)" - fi - - # Error case: bad query - BAD_RESULT=$(python3 -m graphify cypher "INVALID CYPHER" --db "$DB_PATH" 2>&1 | grep -v "^INFO\|^E20" || true) - if echo "$BAD_RESULT" | grep -qi "error\|fail\|traceback"; then - pass "bad cypher query properly errors" - else - fail "bad cypher query did not error" - fi - - # Error case: missing db - MISS_RESULT=$(python3 -m graphify cypher "MATCH (n) RETURN n" --db "/nonexistent/path.db" 2>&1 | grep -v "^INFO\|^E20" || true) - if echo "$MISS_RESULT" | grep -qi "not found\|error"; then - pass "missing db properly errors" - else - fail "missing db did not error" - fi -fi - -echo "" - -# ============================================================ -# TEST 3: Incremental extract (modify a file, re-run) -# ============================================================ -echo "[Test 3] Incremental extract (modify file, re-run)" - -# Modify a file -cat >> "$PROJECT/src/utils.py" << 'PYEOF' - -def new_function(): - """A newly added function.""" - return 42 -PYEOF - -OUTPUT2=$(cd "$PROJECT" && GEMINI_API_KEY=dummy python3 -m graphify extract . --no-semantic --no-cluster 2>&1) - -CLEAN_OUTPUT2=$(echo "$OUTPUT2" | grep -v "^INFO\|^E20") -if echo "$CLEAN_OUTPUT2" | grep -q "incremental"; then - pass "incremental mode detected" -else - skip "incremental mode not detected in output" -fi - -# Check that graph.db still exists and is queryable -if [ -e "$DB_PATH" ]; then - CYPHER_INC=$(python3 -m graphify cypher "MATCH (n:code) RETURN count(n)" --db "$DB_PATH" 2>/dev/null | grep -v "^INFO\|^E20" || echo "ERROR") - if [ "$CYPHER_INC" != "ERROR" ] && [ -n "$CYPHER_INC" ]; then - pass "graph.db still queryable after incremental ($CYPHER_INC nodes)" - else - fail "graph.db not queryable after incremental" - fi -else - skip "graph.db not available for incremental test" -fi - -echo "" - -# ============================================================ -# TEST 4: Python API — storage module direct test -# ============================================================ -echo "[Test 4] Python API — storage module" - -python3 << 'PYTEST' -import sys, json, tempfile, os -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath("__file__")))) -from graphify.storage import init_db, ingest_extraction, ingest_communities, execute_cypher, close_db - -# Small extraction (under 4096 limit) -extraction = { - "nodes": [ - {"id": "fn_main", "label": "main()", "file_type": "code", "source_file": "app.py", "source_location": "L1"}, - {"id": "fn_helper", "label": "helper()", "file_type": "code", "source_file": "app.py", "source_location": "L10"}, - {"id": "concept_arch", "label": "architecture", "file_type": "concept", "source_file": "docs.md"}, - ], - "edges": [ - {"source": "fn_main", "target": "fn_helper", "relation": "calls", "confidence": "EXTRACTED", "source_file": "app.py", "weight": 1.0}, - {"source": "fn_main", "target": "concept_arch", "relation": "implements", "confidence": "INFERRED", "source_file": "app.py", "weight": 0.8}, - ], -} - -d = tempfile.mkdtemp() -db_path = os.path.join(d, "test.db") -db, conn = init_db(db_path) - -# First build -ingest_extraction(conn, extraction, incremental=False) -nodes = execute_cypher(conn, "MATCH (n:code) RETURN count(n)") -assert nodes[0][0] == 2, f"Expected 2 code nodes, got {nodes[0][0]}" - -edges = execute_cypher(conn, "MATCH ()-[e:edge_code_code_calls]->() RETURN count(e)") -assert edges[0][0] == 1, f"Expected 1 code-code-calls edge, got {edges[0][0]}" - -# Communities -ingest_communities(conn, {0: ["fn_main", "fn_helper"], 1: ["concept_arch"]}) -comm = execute_cypher(conn, "MATCH (n:code {id: 'fn_main'}) RETURN n.community") -assert comm[0][0] == 0, f"Expected community 0, got {comm[0][0]}" - -# Incremental update -extraction["nodes"][0]["label"] = "main_v2()" -ingest_extraction(conn, extraction, incremental=True) -updated = execute_cypher(conn, "MATCH (n:code {id: 'fn_main'}) RETURN n.label") -assert updated[0][0] == "main_v2()", f"Expected 'main_v2()', got {updated[0][0]}" - -close_db(db, conn) -print(" All Python API assertions passed") -PYTEST - -if [ $? -eq 0 ]; then - pass "Python API test passed" -else - fail "Python API test failed" -fi - -echo "" - -# ============================================================ -# TEST 5: MCP server tool registration (smoke test) -# ============================================================ -echo "[Test 5] MCP server — cypher_query tool registered" - -python3 << 'MCPTEST' -import sys -# Check that serve.py has cypher_query in its tool list -import importlib.util -spec = importlib.util.find_spec("graphify.serve") -if spec is None: - print(" graphify.serve not found") - sys.exit(1) - -source = open(spec.origin).read() -if "cypher_query" in source and "_tool_cypher_query" in source: - print(" cypher_query tool found in serve.py") - sys.exit(0) -else: - print(" cypher_query tool NOT found in serve.py") - sys.exit(1) -MCPTEST - -if [ $? -eq 0 ]; then - pass "cypher_query tool registered in MCP server" -else - fail "cypher_query tool not found in MCP server" -fi - -echo "" - -# ============================================================ -# TEST 6: Graceful fallback when neug not installed -# ============================================================ -echo "[Test 6] Graceful fallback (simulated)" - -python3 << 'FALLBACK' -import sys, importlib, types - -# Simulate neug not being importable by temporarily removing it -saved = sys.modules.pop("neug", None) -saved_storage = sys.modules.pop("graphify.storage", None) - -# Create a fake module that raises ImportError -blocker = types.ModuleType("neug") -blocker.__spec__ = None - -class NeuGBlocker: - def find_module(self, name, path=None): - if name == "neug" or name.startswith("neug."): - return self - def load_module(self, name): - raise ImportError("simulated: neug not installed") - -sys.meta_path.insert(0, NeuGBlocker()) - -try: - # This should raise ImportError (caught by __main__.py) - from graphify.storage import init_db - print(" ERROR: import should have failed") - sys.exit(1) -except ImportError: - print(" ImportError correctly raised when neug missing") - sys.exit(0) -finally: - sys.meta_path.pop(0) - if saved: - sys.modules["neug"] = saved - if saved_storage: - sys.modules["graphify.storage"] = saved_storage -FALLBACK - -if [ $? -eq 0 ]; then - pass "graceful fallback when neug not installed" -else - fail "fallback test failed" -fi - -echo "" - -# ============================================================ -# Summary -# ============================================================ -echo "======================================" -TOTAL=$((PASS + FAIL + SKIP)) -echo -e " Results: ${GREEN}$PASS passed${NC}, ${RED}$FAIL failed${NC}, ${YELLOW}$SKIP skipped${NC} / $TOTAL total" -echo "======================================" - -if [ $FAIL -gt 0 ]; then - exit 1 -fi -exit 0 diff --git a/tests/test_storage.py b/tests/test_storage.py index c8b21c43dd..ad6bcfd02d 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -119,11 +119,12 @@ def test_ingest_extraction_prune(tmp_db): # --- fallback rel table --- def test_fallback_rel_table(tmp_db): - from graphify.storage import _ensure_rel_table, _created_rel_tables + from graphify.storage import _ensure_rel_table, ensure_schema db, conn = _init(tmp_db) - tbl = _ensure_rel_table(conn, "paper", "document", "cites") + known = ensure_schema(conn) + tbl = _ensure_rel_table(conn, "paper", "document", "cites", known) assert tbl == "edge_paper_document_cites" - assert tbl in _created_rel_tables + assert tbl in known _close(db, conn) From d5c0334eab051a2577640e80c33edd23944f0f7b Mon Sep 17 00:00:00 2001 From: BingqingLyu Date: Fri, 29 May 2026 11:52:26 +0800 Subject: [PATCH 08/10] docs: add NeuG optional extra and cypher CLI examples to README; add storage.py module to ARCHITECTURE.md Co-Authored-By: Claude Opus 4.6 --- ARCHITECTURE.md | 1 + README.md | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 5672bf0df2..a33ec7ceaf 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -26,6 +26,7 @@ Each stage is a single function in its own module. They communicate through plai | `cache.py` | `check_semantic_cache / save_semantic_cache` | files → (cached, uncached) split | | `security.py` | validation helpers | URL / path / label → validated or raises | | `validate.py` | `validate_extraction(data)` | extraction dict → raises on schema errors | +| `storage.py` | `init_db / ingest_extraction / ingest_communities` | extraction dict → NeuG `graph.db` (optional, requires `neug`) | | `serve.py` | `start_server(graph_path)` | graph file path → MCP stdio server | | `watch.py` | `watch(root, flag_path)` | directory → writes flag file on change | | `benchmark.py` | `run_benchmark(graph_path)` | graph file → corpus vs subgraph token comparison | diff --git a/README.md b/README.md index 947d3f7cf6..411fabb465 100644 --- a/README.md +++ b/README.md @@ -156,6 +156,7 @@ Install only what you need: | `video` | Video/audio transcription (faster-whisper + yt-dlp) | `pip install "graphifyy[video]"` | | `mcp` | MCP stdio server | `pip install "graphifyy[mcp]"` | | `neo4j` | Neo4j push support | `pip install "graphifyy[neo4j]"` | +| `neug` | [NeuG](https://github.com/alibaba/neug) embedded graph database — Cypher queries on your graph | `pip install "graphifyy[neug]"` | | `svg` | SVG graph export | `pip install "graphifyy[svg]"` | | `leiden` | Leiden community detection (Python < 3.13 only) | `pip install "graphifyy[leiden]"` | | `ollama` | Ollama local inference | `pip install "graphifyy[ollama]"` | @@ -439,6 +440,9 @@ graphify install # overwrites the skill file /graphify ./raw --graphml # export for Gephi / yEd /graphify ./raw --neo4j # generate cypher.txt for Neo4j /graphify ./raw --neo4j-push bolt://localhost:7687 + +graphify cypher "MATCH (n) RETURN n LIMIT 10" # query graph.db with Cypher (requires neug) +graphify cypher "MATCH (n:code)-[e]->(m) RETURN n.id, e, m.id LIMIT 10" --db path/to/graph.db # default: graphify-out/graph.db /graphify ./raw --watch # auto-sync as files change /graphify ./raw --mcp # start MCP stdio server From c4227b4c27c56d15eed833c7c8e5ebb5091098ee Mon Sep 17 00:00:00 2001 From: BingqingLyu Date: Tue, 9 Jun 2026 12:07:53 +0800 Subject: [PATCH 09/10] perf(storage): use COPY FROM bulk loading for NeuG full builds Replace per-row Cypher CREATE with CSV-based COPY FROM for full builds, achieving 27-69x speedup on node ingestion. Split ingest_extraction into _bulk_ingest (COPY FROM) and _incremental_ingest (per-row MERGE). Co-Authored-By: Claude Opus 4.6 --- graphify/storage.py | 249 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 193 insertions(+), 56 deletions(-) diff --git a/graphify/storage.py b/graphify/storage.py index ce55a8539e..598992c11a 100644 --- a/graphify/storage.py +++ b/graphify/storage.py @@ -11,7 +11,10 @@ """ from __future__ import annotations +import csv +import os import re +import tempfile from pathlib import Path from .build import _FILE_TYPE_SYNONYMS, _normalize_id, _norm_source_file @@ -42,6 +45,18 @@ source_file STRING, community INT64)""", } +_NODE_COLUMNS = { + "code": ["id", "label", "source_file", "source_location", "community"], + "document": ["id", "label", "source_file", "community"], + "paper": ["id", "label", "source_file", "community"], + "image": ["id", "label", "source_file", "community"], + "concept": ["id", "label", "source_file", "community"], + "rationale": ["id", "label", "source_file", "community"], +} + +_EDGE_COLUMNS = ["from_id", "to_id", "relation", "confidence", + "confidence_score", "source_file", "weight"] + # --------------------------------------------------------------------------- # Edge tables — split by (src_type, tgt_type, relation). # --------------------------------------------------------------------------- @@ -74,6 +89,43 @@ def _edge_table_name(src_type: str, tgt_type: str, relation: str) -> str: return f"edge_{src_type}_{tgt_type}_{_sanitize_rel_name(relation)}" +# --------------------------------------------------------------------------- +# CSV helpers for bulk COPY FROM +# --------------------------------------------------------------------------- + +def _sanitize_csv_value(v: object) -> str: + if isinstance(v, str): + return v.replace("\n", "\\n").replace("\r", "") + return str(v) + + +def _write_csv(path: str, rows: list[dict], columns: list[str]) -> int: + if not rows: + return 0 + with open(path, "w", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=columns, extrasaction="ignore", + quoting=csv.QUOTE_ALL) + w.writeheader() + for row in rows: + w.writerow({k: _sanitize_csv_value(row.get(k, "")) for k in columns}) + return len(rows) + + +def _copy_node_csv(conn: object, csv_path: str, table: str) -> None: + conn.execute( + f'COPY {table} FROM "{csv_path}" (header=true, delim=",", escaping=false)' + ) + + +def _copy_rel_csv(conn: object, csv_path: str, tbl: str, + src_table: str, tgt_table: str) -> None: + conn.execute( + f'COPY {tbl} FROM "{csv_path}" ' + f'(from="{src_table}", to="{tgt_table}", ' + f'header=true, delim=",", escaping=false)' + ) + + # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- @@ -135,100 +187,154 @@ def _fix_file_type(ft: str | None) -> str: return ft -def ingest_extraction( +def _bulk_ingest( conn: object, extraction: dict, *, - incremental: bool = False, - prune_sources: list[str] | None = None, - root: str | Path | None = None, + root: str | None = None, known_tables: set[str] | None = None, ) -> dict[str, str]: - """Write an extraction dict into NeuG. + """Full build via COPY FROM — much faster than per-row Cypher CREATE.""" + _known = known_tables if known_tables is not None else set() + nodes = extraction.get("nodes") or [] + edges = extraction.get("edges") or [] - incremental=False: first build — uses CREATE (faster). - incremental=True: update — uses MERGE (upsert). + # --- collect node rows grouped by file_type --- + node_types: dict[str, str] = {} + node_buckets: dict[str, list[dict]] = {ft: [] for ft in _NODE_TABLES} + written_ids: set[str] = set() - Returns node_types dict (id -> file_type) for use by ingest_communities. - """ - _root = str(Path(root).resolve()) if root else None + for node in nodes: + nid = _normalize_id(node.get("id", "")) + if not nid or nid in written_ids: + continue + written_ids.add(nid) + ft = _fix_file_type(node.get("file_type")) + node_types[nid] = ft + row: dict = { + "id": nid, + "label": node.get("label", ""), + "source_file": _norm_source_file(node.get("source_file"), root) or "", + "community": 0, + } + if ft == "code": + row["source_location"] = node.get("source_location") or "" + node_buckets.setdefault(ft, []).append(row) + + # --- collect edge rows grouped by rel table --- + edge_buckets: dict[str, list[dict]] = {} + edge_table_types: dict[str, tuple[str, str]] = {} + + for edge in edges: + src_id = _normalize_id(edge.get("source") or edge.get("from", "")) + tgt_id = _normalize_id(edge.get("target") or edge.get("to", "")) + if not src_id or not tgt_id: + continue + src_ft = node_types.get(src_id) + tgt_ft = node_types.get(tgt_id) + if not src_ft or not tgt_ft: + continue + + rel_raw = edge.get("relation", "") + tbl = _ensure_rel_table(conn, src_ft, tgt_ft, rel_raw, _known) + edge_table_types[tbl] = (src_ft, tgt_ft) + edge_buckets.setdefault(tbl, []).append({ + "from_id": src_id, + "to_id": tgt_id, + "relation": rel_raw, + "confidence": edge.get("confidence", ""), + "confidence_score": float(edge.get("confidence_score", 0.0)), + "source_file": _norm_source_file(edge.get("source_file"), root) or "", + "weight": float(edge.get("weight", 1.0)), + }) + + # --- write CSV + COPY FROM in a temp dir --- + tmp_dir = tempfile.mkdtemp(prefix="graphify_bulk_") + try: + for ft, rows in node_buckets.items(): + if not rows: + continue + csv_path = os.path.join(tmp_dir, f"node_{ft}.csv") + _write_csv(csv_path, rows, _NODE_COLUMNS[ft]) + _copy_node_csv(conn, csv_path, ft) + + for tbl, rows in edge_buckets.items(): + if not rows: + continue + csv_path = os.path.join(tmp_dir, f"edge_{tbl}.csv") + _write_csv(csv_path, rows, _EDGE_COLUMNS) + src_ft, tgt_ft = edge_table_types[tbl] + _copy_rel_csv(conn, csv_path, tbl, src_ft, tgt_ft) + finally: + import shutil + shutil.rmtree(tmp_dir, ignore_errors=True) + + return node_types + + +def _incremental_ingest( + conn: object, + extraction: dict, + *, + prune_sources: list[str] | None = None, + root: str | None = None, + known_tables: set[str] | None = None, +) -> dict[str, str]: + """Incremental update via per-row Cypher MERGE.""" _known = known_tables if known_tables is not None else set() - # --- prune deleted/changed files first --- if prune_sources: for sf in prune_sources: - sf_norm = _norm_source_file(sf, _root) or sf + sf_norm = _norm_source_file(sf, root) or sf for tbl in _NODE_TABLES: conn.execute( f"MATCH (n:{tbl}) WHERE n.source_file = $sf DETACH DELETE n", parameters={"sf": sf_norm}, ) - # --- build node lookup: id -> file_type --- node_types: dict[str, str] = {} nodes = extraction.get("nodes") or [] edges = extraction.get("edges") or [] - # --- write nodes --- - _written_ids: set[str] = set() + written_ids: set[str] = set() _n_errors = 0 for node in nodes: nid = _normalize_id(node.get("id", "")) - if not nid: + if not nid or nid in written_ids: continue + written_ids.add(nid) ft = _fix_file_type(node.get("file_type")) label = node.get("label", "") - sf = _norm_source_file(node.get("source_file"), _root) or "" + sf = _norm_source_file(node.get("source_file"), root) or "" sl = node.get("source_location") or "" node_types[nid] = ft - if nid in _written_ids: - continue - _written_ids.add(nid) try: - if incremental: - if ft == "code": - conn.execute( - f"MERGE (n:code {{id: $nid}}) " - f"ON CREATE SET n.label = $label, " - f"n.source_file = $sf, n.source_location = $sl " - f"ON MATCH SET n.label = $label, " - f"n.source_file = $sf, n.source_location = $sl", - parameters={"nid": nid, "label": label, "sf": sf, "sl": sl}, - ) - else: - conn.execute( - f"MERGE (n:{ft} {{id: $nid}}) " - f"ON CREATE SET n.label = $label, n.source_file = $sf " - f"ON MATCH SET n.label = $label, n.source_file = $sf", - parameters={"nid": nid, "label": label, "sf": sf}, - ) + if ft == "code": + conn.execute( + f"MERGE (n:code {{id: $nid}}) " + f"ON CREATE SET n.label = $label, " + f"n.source_file = $sf, n.source_location = $sl " + f"ON MATCH SET n.label = $label, " + f"n.source_file = $sf, n.source_location = $sl", + parameters={"nid": nid, "label": label, "sf": sf, "sl": sl}, + ) else: - if ft == "code": - conn.execute( - f"CREATE (n:code {{id: $nid, label: $label, " - f"source_file: $sf, source_location: $sl}})", - parameters={"nid": nid, "label": label, "sf": sf, "sl": sl}, - ) - else: - conn.execute( - f"CREATE (n:{ft} {{id: $nid, label: $label, " - f"source_file: $sf}})", - parameters={"nid": nid, "label": label, "sf": sf}, - ) + conn.execute( + f"MERGE (n:{ft} {{id: $nid}}) " + f"ON CREATE SET n.label = $label, n.source_file = $sf " + f"ON MATCH SET n.label = $label, n.source_file = $sf", + parameters={"nid": nid, "label": label, "sf": sf}, + ) except RuntimeError: _n_errors += 1 - # --- write edges --- _e_errors = 0 for edge in edges: - src_key = edge.get("source") or edge.get("from", "") - tgt_key = edge.get("target") or edge.get("to", "") - src_id = _normalize_id(src_key) - tgt_id = _normalize_id(tgt_key) + src_id = _normalize_id(edge.get("source") or edge.get("from", "")) + tgt_id = _normalize_id(edge.get("target") or edge.get("to", "")) if not src_id or not tgt_id: continue - src_ft = node_types.get(src_id) tgt_ft = node_types.get(tgt_id) if not src_ft or not tgt_ft: @@ -238,7 +344,7 @@ def ingest_extraction( conf_raw = edge.get("confidence", "") tbl = _ensure_rel_table(conn, src_ft, tgt_ft, rel_raw, _known) conf_score = float(edge.get("confidence_score", 0.0)) - e_sf = _norm_source_file(edge.get("source_file"), _root) or "" + e_sf = _norm_source_file(edge.get("source_file"), root) or "" weight = float(edge.get("weight", 1.0)) try: @@ -268,6 +374,37 @@ def ingest_extraction( return node_types +def ingest_extraction( + conn: object, + extraction: dict, + *, + incremental: bool = False, + prune_sources: list[str] | None = None, + root: str | Path | None = None, + known_tables: set[str] | None = None, +) -> dict[str, str]: + """Write an extraction dict into NeuG. + + incremental=False: first build — uses COPY FROM bulk loading. + incremental=True: update — uses MERGE (upsert) per row. + + Returns node_types dict (id -> file_type) for use by ingest_communities. + """ + _root = str(Path(root).resolve()) if root else None + + if incremental: + return _incremental_ingest( + conn, extraction, + prune_sources=prune_sources, root=_root, + known_tables=known_tables, + ) + else: + return _bulk_ingest( + conn, extraction, + root=_root, known_tables=known_tables, + ) + + def ingest_communities( conn: object, communities: dict[int, list[str]], From 4f3b2f38db6d26b06f2e1492fc5dccbb8537ff7f Mon Sep 17 00:00:00 2001 From: BingqingLyu Date: Wed, 10 Jun 2026 10:38:01 +0800 Subject: [PATCH 10/10] perf(storage): use DELETE + COPY FROM for incremental updates Replace per-row MERGE with batch DELETE affected source_files + COPY FROM for incremental ingest. Preserves incoming cross-file edges by saving them before deletion and restoring afterwards. Resolves non-delta edge endpoint types from the database before deletion. Benchmark: 1.7-9.4x faster than NetworkX across all repo scales. Correctness verified: node/edge counts match NetworkX after full build, single-file update, multi-file update, and prune operations. Co-Authored-By: Claude Opus 4.6 --- graphify/storage.py | 203 +++++++++++++++++++++++++++++++------------- 1 file changed, 143 insertions(+), 60 deletions(-) diff --git a/graphify/storage.py b/graphify/storage.py index 598992c11a..07aec39afc 100644 --- a/graphify/storage.py +++ b/graphify/storage.py @@ -280,56 +280,130 @@ def _incremental_ingest( root: str | None = None, known_tables: set[str] | None = None, ) -> dict[str, str]: - """Incremental update via per-row Cypher MERGE.""" + """Incremental update via DELETE affected source_files + COPY FROM. + + Much faster than per-row MERGE: deletes nodes whose source_file appears + in the incoming extraction (or in prune_sources), then bulk-inserts the + new data via COPY FROM. Incoming cross-file edges (from unchanged files + into affected nodes) are saved before deletion and restored afterwards. + """ _known = known_tables if known_tables is not None else set() + nodes = extraction.get("nodes") or [] + edges = extraction.get("edges") or [] + # --- collect affected source_files from the incoming data --- + affected_sfs: set[str] = set() if prune_sources: for sf in prune_sources: sf_norm = _norm_source_file(sf, root) or sf - for tbl in _NODE_TABLES: - conn.execute( - f"MATCH (n:{tbl}) WHERE n.source_file = $sf DETACH DELETE n", - parameters={"sf": sf_norm}, - ) + affected_sfs.add(sf_norm) node_types: dict[str, str] = {} - nodes = extraction.get("nodes") or [] - edges = extraction.get("edges") or [] - + node_buckets: dict[str, list[dict]] = {ft: [] for ft in _NODE_TABLES} written_ids: set[str] = set() - _n_errors = 0 + for node in nodes: nid = _normalize_id(node.get("id", "")) if not nid or nid in written_ids: continue written_ids.add(nid) ft = _fix_file_type(node.get("file_type")) - label = node.get("label", "") - sf = _norm_source_file(node.get("source_file"), root) or "" - sl = node.get("source_location") or "" node_types[nid] = ft + sf = _norm_source_file(node.get("source_file"), root) or "" + if sf: + affected_sfs.add(sf) + row: dict = { + "id": nid, + "label": node.get("label", ""), + "source_file": sf, + "community": 0, + } + if ft == "code": + row["source_location"] = node.get("source_location") or "" + node_buckets.setdefault(ft, []).append(row) - try: - if ft == "code": - conn.execute( - f"MERGE (n:code {{id: $nid}}) " - f"ON CREATE SET n.label = $label, " - f"n.source_file = $sf, n.source_location = $sl " - f"ON MATCH SET n.label = $label, " - f"n.source_file = $sf, n.source_location = $sl", - parameters={"nid": nid, "label": label, "sf": sf, "sl": sl}, - ) - else: - conn.execute( - f"MERGE (n:{ft} {{id: $nid}}) " - f"ON CREATE SET n.label = $label, n.source_file = $sf " - f"ON MATCH SET n.label = $label, n.source_file = $sf", - parameters={"nid": nid, "label": label, "sf": sf}, - ) - except RuntimeError: - _n_errors += 1 + # --- resolve types for non-delta edge endpoints (before DELETE) --- + unknown_ids: set[str] = set() + for edge in edges: + for key in ("source", "from", "target", "to"): + eid = _normalize_id(edge.get(key, "")) + if eid and eid not in node_types: + unknown_ids.add(eid) + for nid in unknown_ids: + for tbl in _NODE_TABLES: + try: + rows = list(conn.execute( + f"MATCH (n:{tbl} {{id: $nid}}) RETURN 1", + parameters={"nid": nid}, + )) + if rows: + node_types[nid] = tbl + break + except RuntimeError: + pass + + # --- save incoming cross-file edges before DELETE --- + # Collect IDs of nodes that will be deleted. + affected_node_ids: set[str] = set() + for sf in affected_sfs: + for tbl in _NODE_TABLES: + try: + for row in conn.execute( + f"MATCH (n:{tbl}) WHERE n.source_file = $sf RETURN n.id", + parameters={"sf": sf}, + ): + affected_node_ids.add(row[0]) + except RuntimeError: + pass + + # For each known edge table, find edges where the target is in an affected + # source_file but the source is NOT (incoming from unchanged files). + saved_edge_buckets: dict[str, list[dict]] = {} + saved_edge_types: dict[str, tuple[str, str]] = {} + + for tbl in list(_known): + parts = tbl.split("_", 3) + if len(parts) < 4 or parts[0] != "edge": + continue + src_type, tgt_type = parts[1], parts[2] + + for sf in affected_sfs: + try: + rows = list(conn.execute( + f"MATCH (a:{src_type})-[e:{tbl}]->(b:{tgt_type}) " + f"WHERE b.source_file = $sf " + f"RETURN a.id, b.id, e.relation, e.confidence, " + f"e.confidence_score, e.source_file, e.weight", + parameters={"sf": sf}, + )) + except RuntimeError: + continue + + for row in rows: + if row[0] in affected_node_ids: + continue + saved_edge_types[tbl] = (src_type, tgt_type) + saved_edge_buckets.setdefault(tbl, []).append({ + "from_id": row[0], "to_id": row[1], + "relation": row[2] or "", + "confidence": row[3] or "", + "confidence_score": float(row[4] or 0.0), + "source_file": row[5] or "", + "weight": float(row[6] or 1.0), + }) + + # --- DELETE nodes from affected source_files --- + for sf in affected_sfs: + for tbl in _NODE_TABLES: + conn.execute( + f"MATCH (n:{tbl}) WHERE n.source_file = $sf DETACH DELETE n", + parameters={"sf": sf}, + ) + + # --- collect delta edge rows --- + edge_buckets: dict[str, list[dict]] = {} + edge_table_types: dict[str, tuple[str, str]] = {} - _e_errors = 0 for edge in edges: src_id = _normalize_id(edge.get("source") or edge.get("from", "")) tgt_id = _normalize_id(edge.get("target") or edge.get("to", "")) @@ -341,35 +415,44 @@ def _incremental_ingest( continue rel_raw = edge.get("relation", "") - conf_raw = edge.get("confidence", "") tbl = _ensure_rel_table(conn, src_ft, tgt_ft, rel_raw, _known) - conf_score = float(edge.get("confidence_score", 0.0)) - e_sf = _norm_source_file(edge.get("source_file"), root) or "" - weight = float(edge.get("weight", 1.0)) + edge_table_types[tbl] = (src_ft, tgt_ft) + edge_buckets.setdefault(tbl, []).append({ + "from_id": src_id, + "to_id": tgt_id, + "relation": rel_raw, + "confidence": edge.get("confidence", ""), + "confidence_score": float(edge.get("confidence_score", 0.0)), + "source_file": _norm_source_file(edge.get("source_file"), root) or "", + "weight": float(edge.get("weight", 1.0)), + }) - try: - conn.execute( - f"MATCH (a:{src_ft} {{id: $src_id}}), " - f"(b:{tgt_ft} {{id: $tgt_id}}) " - f"CREATE (a)-[:{tbl} {{relation: $rel, confidence: $conf, " - f"confidence_score: $conf_score, source_file: $e_sf, " - f"weight: $weight}}]->(b)", - parameters={ - "src_id": src_id, "tgt_id": tgt_id, - "rel": rel_raw, "conf": conf_raw, - "conf_score": conf_score, "e_sf": e_sf, - "weight": weight, - }, - ) - except RuntimeError: - _e_errors += 1 - - if _n_errors or _e_errors: - import logging - logging.getLogger(__name__).warning( - "NeuG ingest: %d node(s) and %d edge(s) skipped due to errors", - _n_errors, _e_errors, - ) + # --- merge saved incoming edges back --- + for tbl, rows in saved_edge_buckets.items(): + edge_buckets.setdefault(tbl, []).extend(rows) + if tbl not in edge_table_types: + edge_table_types[tbl] = saved_edge_types[tbl] + + # --- COPY FROM bulk insert --- + tmp_dir = tempfile.mkdtemp(prefix="graphify_inc_") + try: + for ft, rows in node_buckets.items(): + if not rows: + continue + csv_path = os.path.join(tmp_dir, f"node_{ft}.csv") + _write_csv(csv_path, rows, _NODE_COLUMNS[ft]) + _copy_node_csv(conn, csv_path, ft) + + for tbl, rows in edge_buckets.items(): + if not rows: + continue + csv_path = os.path.join(tmp_dir, f"edge_{tbl}.csv") + _write_csv(csv_path, rows, _EDGE_COLUMNS) + src_ft, tgt_ft = edge_table_types[tbl] + _copy_rel_csv(conn, csv_path, tbl, src_ft, tgt_ft) + finally: + import shutil + shutil.rmtree(tmp_dir, ignore_errors=True) return node_types