From 49df4663856311dcaab909d052d55d1200aad470 Mon Sep 17 00:00:00 2001 From: safishamsi Date: Wed, 15 Jul 2026 10:31:59 +0100 Subject: [PATCH] fix(sql): recover CREATE FUNCTION/PROCEDURE from tree-sitter ERROR nodes (#1910) tree-sitter-sql parses PL/pgSQL CREATE FUNCTION statements (OUT/INOUT params, tagged dollar quotes, PERFORM/:= body statements) as ERROR nodes, and the dispatch loop had no branch for them, so the functions were silently dropped from the graph. Handle ERROR nodes inside walk() (they can nest inside a merged create_function during multi-statement error recovery) and dispatch top-level ERROR statements to it. The branch regex-scans the raw node text for every CREATE [OR REPLACE] FUNCTION/PROCEDURE, mirroring the existing fb_proc_or_trigger and has_error fallbacks, with a name class that keeps schema-qualified names (exposed.important_function) whole. The PL/pgSQL body is not scanned for FROM/JOIN references to avoid junk reads_from targets, and node ids match the clean create_function branch so seen_ids dedups consistently. Co-Authored-By: Claude Opus 4.8 (1M context) --- graphify/extractors/sql.py | 21 ++++++++++++++++++++- tests/fixtures/sample_plpgsql.sql | 27 +++++++++++++++++++++++++++ tests/test_multilang.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 76 insertions(+), 1 deletion(-) create mode 100644 tests/fixtures/sample_plpgsql.sql diff --git a/graphify/extractors/sql.py b/graphify/extractors/sql.py index fa87982f..b9ce6fd8 100644 --- a/graphify/extractors/sql.py +++ b/graphify/extractors/sql.py @@ -191,6 +191,25 @@ def extract_sql(path: Path, content: str | bytes | None = None) -> dict: tbl_nid = table_nids.get(tbl_name.lower()) or _make_id(stem, tbl_name) _add_edge(trig_nid, tbl_nid, "triggers", line) + elif t == "ERROR": + # tree-sitter-sql cannot parse PL/pgSQL CREATE FUNCTION/PROCEDURE + # bodies (OUT/INOUT params, tagged dollar quotes, PERFORM, :=) and + # emits an ERROR node instead, silently dropping the object. + # Regex-scan the raw text as fallback, mirroring the + # fb_proc_or_trigger recovery below. One ERROR blob can swallow + # several statements, so scan for every CREATE in it. We deliberately + # do not scan the body for FROM/JOIN references: PL/pgSQL loop + # variables and locals would produce junk reads_from targets. + text = _read(node) + for m in re.finditer( + r"CREATE\s+(?:OR\s+REPLACE\s+)?(?:FUNCTION|PROCEDURE)\s+([\w$.]+)", + text, re.IGNORECASE, + ): + name = m.group(1) + m_line = line + text[: m.start()].count("\n") + nid = _make_id(stem, name) + _add_node(nid, f"{name}()", m_line) + elif t == "fb_proc_or_trigger": text = _read(node) m = re.match( @@ -249,7 +268,7 @@ def extract_sql(path: Path, content: str | bytes | None = None) -> dict: if stmt.type == "statement": for child in stmt.children: walk(child) - elif stmt.type in ("fb_proc_or_trigger", "set_term", "declare_external_function"): + elif stmt.type in ("fb_proc_or_trigger", "set_term", "declare_external_function", "ERROR"): walk(stmt) # Global regex fallback: catch any REFERENCES missed due to ERROR nodes in the parse tree diff --git a/tests/fixtures/sample_plpgsql.sql b/tests/fixtures/sample_plpgsql.sql new file mode 100644 index 00000000..de2ef161 --- /dev/null +++ b/tests/fixtures/sample_plpgsql.sql @@ -0,0 +1,27 @@ +CREATE TABLE accounts ( + id INT PRIMARY KEY, + name TEXT +); + +CREATE FUNCTION exposed.important_function(a int, OUT fuzz text) LANGUAGE plpgsql AS $$ +BEGIN + PERFORM 1; + fuzz := 'x'; +END +$$; + +CREATE OR REPLACE FUNCTION tagged_quote_fn(n int) RETURNS int LANGUAGE plpgsql AS $fn$ +DECLARE + total int := 0; +BEGIN + total := total + n; + RETURN total; +END +$fn$; + +CREATE FUNCTION plain_sql_fn() RETURNS int LANGUAGE sql AS 'SELECT 1'; + +CREATE TABLE audit_log ( + id INT PRIMARY KEY, + account_id INT REFERENCES accounts +); diff --git a/tests/test_multilang.py b/tests/test_multilang.py index 5ad78f72..0cc26539 100644 --- a/tests/test_multilang.py +++ b/tests/test_multilang.py @@ -513,3 +513,32 @@ def test_sql_schema_qualified_alter_fk(): for e in fk_edges: assert e["source"] in node_ids, f"dangling source: {e['source']}" assert e["target"] in node_ids, f"dangling target: {e['target']}" + +def test_sql_plpgsql_functions_survive_parse_errors(): + """PL/pgSQL bodies make tree-sitter-sql emit ERROR nodes; the functions + must still be extracted (#1910), without cascading into later statements.""" + r = _extract_sql_or_skip("sample_plpgsql.sql") + labels = [n["label"] for n in r["nodes"]] + # Both PL/pgSQL functions extracted, schema-qualified name kept whole + assert "exposed.important_function()" in labels + assert "tagged_quote_fn()" in labels + # Tables before and after the broken functions still extract + assert any("accounts" in l for l in labels) + assert any("audit_log" in l for l in labels) + # No spurious or empty nodes from the error recovery + for l in labels: + assert l, "empty node label" + assert l != "ERROR" + # Every function got a contains edge from the file node + contains_targets = {e["target"] for e in r["edges"] if e["relation"] == "contains"} + fn_ids = {n["id"] for n in r["nodes"] if n["label"].endswith("()")} + assert fn_ids <= contains_targets + +def test_sql_plpgsql_clean_function_not_double_emitted(): + """A cleanly-parsed LANGUAGE sql function in the same file is emitted once.""" + r = _extract_sql_or_skip("sample_plpgsql.sql") + labels = [n["label"] for n in r["nodes"]] + assert labels.count("plain_sql_fn()") == 1 + # And nothing else is duplicated either + ids = [n["id"] for n in r["nodes"]] + assert len(ids) == len(set(ids))