fix(pg_introspect): emit FK DDL before function stubs so unparseable routines can't drop references edges (#1854)

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

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."""