From 5d0392524a77ae7cd7e50bf0c954705e1bd99bff Mon Sep 17 00:00:00 2001 From: zuyeyang <47495598+zuyeyang@users.noreply.github.com> Date: Fri, 8 May 2026 12:31:25 -0700 Subject: [PATCH] fix(extract): add ALTER TABLE FK extraction + schema-qualified name support for SQL (#779) The SQL parser (`extract_sql`) previously only extracted foreign key relationships defined inline within CREATE TABLE column definitions. FK constraints added via ALTER TABLE ... ADD CONSTRAINT ... FOREIGN KEY ... REFERENCES were silently ignored. Additionally, `_obj_name()` only read the first identifier child of object_reference nodes, so schema-qualified names like `Sales.Customer` were truncated to just `Sales`. Changes: - Add `alter_table` handler to `walk()` that extracts FK edges from ALTER TABLE ... ADD CONSTRAINT ... FOREIGN KEY ... REFERENCES - Fix `_obj_name()` to read the full object_reference text, preserving schema-qualified names (e.g. `Sales.Customer`) - Fix inline FK resolution in create_table and _walk_from_refs to use full object_reference text instead of first identifier only --- graphify/extract.py | 45 ++++++++++++++++------ tests/fixtures/sample_alter_fk.sql | 12 ++++++ tests/fixtures/sample_schema_qualified.sql | 11 ++++++ tests/test_multilang.py | 27 +++++++++++++ 4 files changed, 83 insertions(+), 12 deletions(-) create mode 100644 tests/fixtures/sample_alter_fk.sql create mode 100644 tests/fixtures/sample_schema_qualified.sql diff --git a/graphify/extract.py b/graphify/extract.py index c0dc16a..6000e02 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -2505,9 +2505,7 @@ def extract_sql(path: Path) -> dict: def _obj_name(n) -> str | None: for c in n.children: if c.type == "object_reference": - for cc in c.children: - if cc.type == "identifier": - return _read(cc) + return _read(c) return None def _add_node(nid: str, label: str, line: int) -> None: @@ -2546,9 +2544,7 @@ def extract_sql(path: Path) -> dict: if cc.type == "keyword_references": found_ref = True elif found_ref and cc.type == "object_reference": - for ccc in cc.children: - if ccc.type == "identifier": - ref_name = _read(ccc) + ref_name = _read(cc) break if ref_name: ref_nid = _make_id(stem, ref_name) @@ -2577,6 +2573,33 @@ def extract_sql(path: Path) -> dict: _add_node(nid, f"{name}()", line) _walk_from_refs(node, nid, line) + elif t == "alter_table": + name = _obj_name(node) + if name: + src_nid = table_nids.get(name.lower()) + if not src_nid: + src_nid = _make_id(stem, name) + _add_node(src_nid, name, line) + table_nids[name.lower()] = src_nid + for child in node.children: + if child.type == "add_constraint": + for cc in child.children: + if cc.type != "constraint": + continue + found_ref = False + ref_name: str | None = None + for ccc in cc.children: + if ccc.type == "keyword_references": + found_ref = True + elif found_ref and ccc.type == "object_reference": + ref_name = _read(ccc) + break + if ref_name: + ref_nid = table_nids.get(ref_name.lower()) + if not ref_nid: + ref_nid = _make_id(stem, ref_name) + _add_edge(src_nid, ref_nid, "references", line) + for child in node.children: walk(child) @@ -2587,12 +2610,10 @@ def extract_sql(path: Path) -> dict: if c.type == "relation": for cc in c.children: if cc.type == "object_reference": - for ccc in cc.children: - if ccc.type == "identifier": - tbl = _read(ccc) - tbl_nid = _make_id(stem, tbl) - _add_edge(caller_nid, tbl_nid, "reads_from", - c.start_point[0] + 1) + tbl = _read(cc) + tbl_nid = _make_id(stem, tbl) + _add_edge(caller_nid, tbl_nid, "reads_from", + c.start_point[0] + 1) for child in node.children: _walk_from_refs(child, caller_nid, line) diff --git a/tests/fixtures/sample_alter_fk.sql b/tests/fixtures/sample_alter_fk.sql new file mode 100644 index 0000000..60d884a --- /dev/null +++ b/tests/fixtures/sample_alter_fk.sql @@ -0,0 +1,12 @@ +CREATE TABLE customers ( + id SERIAL PRIMARY KEY, + name TEXT NOT NULL +); + +CREATE TABLE orders ( + id SERIAL PRIMARY KEY, + customer_id INT, + total NUMERIC +); + +ALTER TABLE orders ADD CONSTRAINT fk_customer FOREIGN KEY (customer_id) REFERENCES customers(id); diff --git a/tests/fixtures/sample_schema_qualified.sql b/tests/fixtures/sample_schema_qualified.sql new file mode 100644 index 0000000..c8cf0f0 --- /dev/null +++ b/tests/fixtures/sample_schema_qualified.sql @@ -0,0 +1,11 @@ +CREATE TABLE Sales.Customer ( + CustomerID SERIAL PRIMARY KEY, + Name TEXT NOT NULL +); + +CREATE TABLE Sales.SalesOrder ( + OrderID SERIAL PRIMARY KEY, + CustomerID INT REFERENCES Sales.Customer(CustomerID) +); + +ALTER TABLE Sales.SalesOrder ADD CONSTRAINT fk_cust FOREIGN KEY (CustomerID) REFERENCES Sales.Customer(CustomerID); diff --git a/tests/test_multilang.py b/tests/test_multilang.py index e462003..637dcd1 100644 --- a/tests/test_multilang.py +++ b/tests/test_multilang.py @@ -252,3 +252,30 @@ def test_sql_no_dangling_edges(): node_ids = {n["id"] for n in r["nodes"]} for e in r["edges"]: assert e["source"] in node_ids, f"dangling source: {e['source']}" + +def test_sql_alter_table_fk_edge(): + """ALTER TABLE ... FOREIGN KEY ... REFERENCES produces a references edge.""" + r = extract_sql(FIXTURES / "sample_alter_fk.sql") + fk_edges = [e for e in r["edges"] if e["relation"] == "references"] + assert len(fk_edges) >= 1 + node_ids = {n["id"] for n in r["nodes"]} + 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_schema_qualified_names(): + """Schema-qualified table names (Schema.Table) are preserved.""" + r = extract_sql(FIXTURES / "sample_schema_qualified.sql") + labels = [n["label"] for n in r["nodes"]] + assert any("Sales.Customer" in l for l in labels) + assert any("Sales.SalesOrder" in l for l in labels) + +def test_sql_schema_qualified_alter_fk(): + """ALTER TABLE with schema-qualified names produces correct edges.""" + r = extract_sql(FIXTURES / "sample_schema_qualified.sql") + fk_edges = [e for e in r["edges"] if e["relation"] == "references"] + assert len(fk_edges) >= 1 + node_ids = {n["id"] for n in r["nodes"]} + 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']}"