From 4bb0b6f2f14ecfdb7dac9a79899f4276e1adc338 Mon Sep 17 00:00:00 2001 From: sekmur Date: Wed, 15 Jul 2026 11:00:19 -0700 Subject: [PATCH] fix(pg_introspect): emit FK DDL before function stubs so unparseable routines can't drop references edges (#1854) A C-language routine's body from information_schema.routines is just the C symbol name, so its reconstructed stub (CREATE FUNCTION ... AS $gfx$ name $gfx$ LANGUAGE c) is unparseable by tree-sitter-sql, and the parser's error recovery consumes the statements that follow. With FK ALTER TABLEs emitted last, every references edge was silently lost on any DB with a common extension installed (uuid-ossp, pgcrypto, pg_trgm, fuzzystrmatch, ...). Emit the FK block before the routine DDL. FK statements only reference tables, which are emitted first, so the order is always safe. On a real DB with 35 tables / 41 FKs / 41 extension routines: 41/41 references edges vs 1/41 before. Co-Authored-By: Claude Fable 5 --- graphify/pg_introspect.py | 26 ++++++++++++-------- tests/test_pg_introspect.py | 48 +++++++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 10 deletions(-) diff --git a/graphify/pg_introspect.py b/graphify/pg_introspect.py index 2a2cca628..d183c0512 100644 --- a/graphify/pg_introspect.py +++ b/graphify/pg_introspect.py @@ -113,6 +113,22 @@ def introspect_postgres(dsn: str | None = None) -> dict: else: ddl.append(f"CREATE VIEW {_quote_ident(schema)}.{_quote_ident(name)} AS SELECT 1;") + # FK edges — one ALTER TABLE per constraint (handles composite FKs correctly). + # Emitted BEFORE the function DDL: routine bodies the grammar can't parse + # (notably C-language extension functions, whose "body" is just the C symbol + # name) put tree-sitter into error recovery that consumes the statements + # after them — with FKs last, every 'references' edge is silently lost on + # any DB with a common extension installed (#1854). FK statements only + # reference tables, which are emitted first, so this order is always safe. + for constraint_name, t_schema, t_name, cols, r_schema, r_name, r_cols in fks: + col_list = ", ".join(_quote_ident(c) for c in cols) + ref_col_list = ", ".join(_quote_ident(c) for c in r_cols) + ddl.append( + f"ALTER TABLE {_quote_ident(t_schema)}.{_quote_ident(t_name)} " + f"ADD CONSTRAINT {_quote_ident(constraint_name)} " + f"FOREIGN KEY ({col_list}) REFERENCES {_quote_ident(r_schema)}.{_quote_ident(r_name)}({ref_col_list});" + ) + # Functions & Procedures — real body if available, stub if NULL # Use $gfx$ as the dollar-quote tag to avoid collision with $$ inside bodies. # Use external_language from the catalog; fall back to plpgsql if NULL/blank. @@ -128,16 +144,6 @@ def introspect_postgres(dsn: str | None = None) -> dict: f" AS $gfx$ {actual_body} $gfx$ LANGUAGE {lang};" ) - # FK edges — one ALTER TABLE per constraint (handles composite FKs correctly) - for constraint_name, t_schema, t_name, cols, r_schema, r_name, r_cols in fks: - col_list = ", ".join(_quote_ident(c) for c in cols) - ref_col_list = ", ".join(_quote_ident(c) for c in r_cols) - ddl.append( - f"ALTER TABLE {_quote_ident(t_schema)}.{_quote_ident(t_name)} " - f"ADD CONSTRAINT {_quote_ident(constraint_name)} " - f"FOREIGN KEY ({col_list}) REFERENCES {_quote_ident(r_schema)}.{_quote_ident(r_name)}({ref_col_list});" - ) - ddl_string = "\n".join(ddl) # Determine host/dbname for virtual path DSN sanitization diff --git a/tests/test_pg_introspect.py b/tests/test_pg_introspect.py index fb4749b3c..e29e4374b 100644 --- a/tests/test_pg_introspect.py +++ b/tests/test_pg_introspect.py @@ -283,6 +283,54 @@ def test_pg_introspect_fk_query_avoids_privilege_filtered_view(): assert len(ref_edges) == 1 +def test_pg_introspect_fk_edges_survive_unparseable_function_stubs(): + """#1854: FK edges must survive routines whose reconstructed DDL the SQL + grammar cannot parse. + + A C-language routine's "body" from information_schema.routines is just the + C symbol name, so its reconstructed stub — + CREATE FUNCTION "public"."levenshtein"() RETURNS void + AS $gfx$ levenshtein $gfx$ LANGUAGE c; + — is unparseable, and tree-sitter's error recovery consumes the statements + that follow it. With FK ALTER TABLEs emitted after the function DDL, every + 'references' edge was silently lost on any DB with a common extension + installed (uuid-ossp, pgcrypto, pg_trgm, fuzzystrmatch, …). FKs must be + emitted before the routine DDL so an unparseable stub can only damage + what follows it.""" + n = 6 + mock_tables = [("public", f"t{i}", "BASE TABLE") for i in range(n + 1)] + mock_views = [] + # One C-language extension routine (unparseable stub) + one plpgsql routine + mock_routines = [ + ("public", "levenshtein", "FUNCTION", "levenshtein", "C"), + ("public", "trigfunc", "FUNCTION", "BEGIN SELECT 1; END;", "PLPGSQL"), + ] + mock_fks = [ + (f"fk{i}", "public", f"t{i}", ["ref_id"], "public", f"t{i+1}", ["id"]) + for i in range(n) + ] + + mock_psycopg = _make_mock_psycopg(mock_tables, mock_views, mock_routines, mock_fks) + + with patch.dict("sys.modules", {"psycopg": mock_psycopg}): + res = introspect_postgres("postgresql://myuser:secret@myhost/mydb") + + errors = validate_extraction(res) + assert errors == [], f"Validation errors: {errors}" + + ids_to_labels = {node["id"]: node["label"] for node in res["nodes"]} + ref_pairs = { + (ids_to_labels[e["source"]], ids_to_labels[e["target"]]) + for e in res["edges"] + if e["relation"] == "references" + } + expected = {(_q("public", f"t{i}"), _q("public", f"t{i+1}")) for i in range(n)} + assert ref_pairs == expected, ( + f"FK edges lost to parser error recovery: expected {n}, " + f"got {len(ref_pairs)}: {sorted(ref_pairs)}" + ) + + def test_pg_introspect_connection_error(): """A psycopg.OperationalError must be re-raised as ConnectionError with a sanitized message (no DSN/credentials) and no stack-trace noise."""