mirror of
https://github.com/safishamsi/graphify.git
synced 2026-09-09 23:25:47 +00:00
Azeem1985
62371ad28d
The postgres introspection path synthesizes DDL and parses it with the SQL grammar, but tree-sitter-sql was not in the postgres extra, so `pip install graphifyy[postgres]` left it unable to parse and returned a silent empty result. Add tree-sitter-sql to the extra and raise an actionable ImportError (with the install hint) when the grammar is missing or has an incompatible ABI, instead of reporting zero nodes.
402 lines
16 KiB
Python
402 lines
16 KiB
Python
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
import pytest
|
|
from pathlib import Path
|
|
|
|
pytest.importorskip("tree_sitter_sql", reason="tree-sitter-sql not installed; skip pg_introspect tests")
|
|
|
|
from graphify.pg_introspect import introspect_postgres
|
|
from graphify.validate import validate_extraction
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shared mock infrastructure
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _make_mock_psycopg(tables, views, routines, fks,
|
|
host="myhost", dbname="mydb",
|
|
connect_raises=None):
|
|
"""Return a mock psycopg module wired to the provided catalog data.
|
|
|
|
``routines`` rows must be 5-tuples: (schema, name, rtype, body, ext_lang).
|
|
``fks`` rows must be 7-tuples:
|
|
(constraint_name, t_schema, t_name, [cols], r_schema, r_name, [r_cols])
|
|
``connect_raises``, if set, is an exception *instance* raised by connect().
|
|
"""
|
|
|
|
executed_queries: list[str] = []
|
|
|
|
class MockCursor:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
def execute(self, query, params=None):
|
|
self.query = query
|
|
executed_queries.append(query)
|
|
|
|
def fetchall(self):
|
|
q = self.query.strip().lower()
|
|
if "information_schema.tables" in q:
|
|
return tables
|
|
elif "information_schema.views" in q:
|
|
return views
|
|
elif "information_schema.routines" in q:
|
|
return routines
|
|
elif "pg_constraint" in q:
|
|
return fks
|
|
return []
|
|
|
|
class MockConnection:
|
|
def execute(self, query):
|
|
pass
|
|
|
|
def cursor(self):
|
|
return MockCursor()
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
@property
|
|
def info(self):
|
|
info_mock = MagicMock()
|
|
info_mock.dsn = f"host={host} dbname={dbname} user=myuser password=secret"
|
|
return info_mock
|
|
|
|
mock_psycopg = MagicMock()
|
|
if connect_raises is not None:
|
|
mock_psycopg.connect.side_effect = connect_raises
|
|
# Make the exception type available as an attribute so the module can
|
|
# reference psycopg.OperationalError in the except clause.
|
|
mock_psycopg.OperationalError = type(connect_raises)
|
|
else:
|
|
mock_psycopg.connect.return_value = MockConnection()
|
|
mock_psycopg.OperationalError = Exception # unused path but must exist
|
|
mock_psycopg.conninfo.conninfo_to_dict.return_value = {
|
|
"host": host,
|
|
"dbname": dbname,
|
|
}
|
|
mock_psycopg._executed_queries = executed_queries
|
|
return mock_psycopg
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _q(schema: str, name: str) -> str:
|
|
"""Return the label form that tree-sitter produces for a quoted identifier.
|
|
|
|
pg_introspect emits CREATE TABLE "schema"."name" — tree-sitter reads the
|
|
object_reference text verbatim (quotes included), so the node label is
|
|
'"schema"."name"', not 'schema.name'.
|
|
"""
|
|
return f'"{schema}"."{name}"'
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_pg_introspect_success():
|
|
"""Baseline: tables, views, routines, and a single-column FK all survive."""
|
|
mock_tables = [
|
|
("public", "users", "BASE TABLE"),
|
|
("public", "orders", "BASE TABLE"),
|
|
]
|
|
mock_views = [
|
|
("public", "active_users", "SELECT * FROM public.users WHERE active = true"),
|
|
]
|
|
# 5-tuple: schema, name, rtype, body, ext_lang
|
|
mock_routines = [
|
|
("public", "calculate_total", "FUNCTION", "SELECT 42;", "SQL"),
|
|
("public", "do_nothing", "PROCEDURE", None, "PLPGSQL"),
|
|
]
|
|
# 7-tuple: constraint_name, t_schema, t_name, cols[], r_schema, r_name, r_cols[]
|
|
mock_fks = [
|
|
("fk_orders_user_id", "public", "orders", ["user_id"], "public", "users", ["id"]),
|
|
]
|
|
|
|
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:mypassword@myhost/mydb")
|
|
|
|
# 1. validate_extraction must pass
|
|
errors = validate_extraction(res)
|
|
assert errors == [], f"Validation errors: {errors}"
|
|
|
|
# 2. source_file must be the sanitized virtual path (no credentials)
|
|
expected_source = "postgresql:/myhost/mydb"
|
|
for node in res["nodes"]:
|
|
assert node["source_file"] == expected_source
|
|
for edge in res["edges"]:
|
|
assert edge["source_file"] == expected_source
|
|
|
|
# 3. Expected node labels. pg_introspect double-quotes identifiers in DDL,
|
|
# so tree-sitter returns the raw quoted text as the object_reference.
|
|
node_labels = {n["label"] for n in res["nodes"]}
|
|
assert _q("public", "users") in node_labels, f"users missing; got {node_labels}"
|
|
assert _q("public", "orders") in node_labels, f"orders missing; got {node_labels}"
|
|
# Views keep the schema-qualified label (quoted schema, unquoted body)
|
|
assert _q("public", "active_users") in node_labels, f"active_users missing; got {node_labels}"
|
|
# Functions: label is "<quoted-sig>()"
|
|
assert f'{_q("public", "calculate_total")}()' in node_labels, f"calculate_total() missing; got {node_labels}"
|
|
assert f'{_q("public", "do_nothing")}()' in node_labels, f"do_nothing() missing; got {node_labels}"
|
|
|
|
# 4. File node (label = dbname)
|
|
file_nodes = [n for n in res["nodes"] if n["file_type"] == "code" and n["label"] == "mydb"]
|
|
assert len(file_nodes) == 1
|
|
|
|
# 5. FK references edge: orders → users, exactly once
|
|
users_nid = next(n["id"] for n in res["nodes"] if n["label"] == _q("public", "users"))
|
|
orders_nid = next(n["id"] for n in res["nodes"] if n["label"] == _q("public", "orders"))
|
|
ref_edges = [
|
|
e for e in res["edges"]
|
|
if e["source"] == orders_nid and e["target"] == users_nid and e["relation"] == "references"
|
|
]
|
|
assert len(ref_edges) == 1, f"Expected exactly 1 references edge, got {len(ref_edges)}"
|
|
|
|
|
|
def test_pg_introspect_quoted_identifiers():
|
|
"""Reserved-word and special-character table names must survive DDL round-trip.
|
|
|
|
'order' is a SQL reserved word; 'user-data' contains a hyphen — both would
|
|
produce invalid DDL without quoting, causing tree-sitter to silently drop
|
|
those tables and any FK touching them.
|
|
"""
|
|
mock_tables = [
|
|
("public", "order", "BASE TABLE"), # reserved word
|
|
("public", "user-data", "BASE TABLE"), # hyphen
|
|
]
|
|
mock_views = []
|
|
mock_routines = []
|
|
# FK: user-data.owner_id → order.id
|
|
mock_fks = [
|
|
("fk_userdata_order", "public", "user-data", ["owner_id"], "public", "order", ["id"]),
|
|
]
|
|
|
|
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}"
|
|
|
|
node_labels = {n["label"] for n in res["nodes"]}
|
|
|
|
# Both tables must appear as nodes (quoted form expected from tree-sitter)
|
|
assert _q("public", "order") in node_labels, \
|
|
f"'order' table missing; labels={node_labels}"
|
|
assert _q("public", "user-data") in node_labels, \
|
|
f"'user-data' table missing; labels={node_labels}"
|
|
|
|
# FK references edge must exist
|
|
ref_edges = [e for e in res["edges"] if e["relation"] == "references"]
|
|
assert len(ref_edges) >= 1, "Expected at least one references edge for the FK"
|
|
|
|
|
|
def test_pg_introspect_composite_fk():
|
|
"""A 2-column composite FK must produce exactly ONE references edge, not two.
|
|
|
|
The old code emitted one ADD CONSTRAINT per row of the FK query (one row
|
|
per key column), causing duplicate edges for composite keys.
|
|
"""
|
|
mock_tables = [
|
|
("public", "products", "BASE TABLE"),
|
|
("public", "order_items", "BASE TABLE"),
|
|
]
|
|
mock_views = []
|
|
mock_routines = []
|
|
# Single composite FK: order_items(order_id, product_id) → products(order_id, product_id)
|
|
mock_fks = [
|
|
(
|
|
"fk_order_items_composite",
|
|
"public", "order_items",
|
|
["order_id", "product_id"],
|
|
"public", "products",
|
|
["order_id", "product_id"],
|
|
),
|
|
]
|
|
|
|
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}"
|
|
|
|
products_nid = next(
|
|
n["id"] for n in res["nodes"] if n["label"] == _q("public", "products")
|
|
)
|
|
order_items_nid = next(
|
|
n["id"] for n in res["nodes"] if n["label"] == _q("public", "order_items")
|
|
)
|
|
|
|
ref_edges = [
|
|
e for e in res["edges"]
|
|
if e["source"] == order_items_nid
|
|
and e["target"] == products_nid
|
|
and e["relation"] == "references"
|
|
]
|
|
assert len(ref_edges) == 1, (
|
|
f"Expected exactly 1 references edge for composite FK, got {len(ref_edges)}"
|
|
)
|
|
|
|
|
|
def test_pg_introspect_fk_query_avoids_privilege_filtered_view():
|
|
"""#1746: information_schema.referential_constraints only shows constraints
|
|
where the current user has WRITE access to the referencing table (owner or
|
|
a privilege other than SELECT). A read-only introspection role therefore
|
|
gets zero FK rows — while tables/views/routines all still appear, since
|
|
SELECT is enough for those views — and the graph silently loses every
|
|
'references' edge. The FK query must read pg_catalog.pg_constraint, which
|
|
is not privilege-filtered."""
|
|
mock_tables = [
|
|
("public", "users", "BASE TABLE"),
|
|
("public", "orders", "BASE TABLE"),
|
|
]
|
|
mock_fks = [
|
|
("fk_orders_user_id", "public", "orders", ["user_id"], "public", "users", ["id"]),
|
|
]
|
|
|
|
mock_psycopg = _make_mock_psycopg(mock_tables, [], [], mock_fks)
|
|
|
|
with patch.dict("sys.modules", {"psycopg": mock_psycopg}):
|
|
res = introspect_postgres("postgresql://readonly:secret@myhost/mydb")
|
|
|
|
constraint_queries = [
|
|
q for q in mock_psycopg._executed_queries if "constraint" in q.lower()
|
|
]
|
|
assert constraint_queries, "no FK query was executed"
|
|
assert all(
|
|
"referential_constraints" not in q.lower() for q in constraint_queries
|
|
), "FK query must not read information_schema.referential_constraints (privilege-filtered, #1746)"
|
|
assert any("pg_constraint" in q.lower() for q in constraint_queries)
|
|
|
|
# And the FK still becomes a references edge end-to-end
|
|
ref_edges = [e for e in res["edges"] if e["relation"] == "references"]
|
|
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."""
|
|
|
|
class FakeOperationalError(Exception):
|
|
pass
|
|
|
|
raw_error = FakeOperationalError(
|
|
'connection to server at "myhost" (127.0.0.1), port 5432 failed: '
|
|
'FATAL: password authentication failed for user "myuser"\n'
|
|
"DETAIL: Connection matched pg_hba.conf line 1: …"
|
|
)
|
|
|
|
mock_psycopg = _make_mock_psycopg([], [], [], [], connect_raises=raw_error)
|
|
|
|
with patch.dict("sys.modules", {"psycopg": mock_psycopg}):
|
|
with pytest.raises(ConnectionError) as exc_info:
|
|
introspect_postgres("postgresql://myuser:secret@myhost/mydb")
|
|
|
|
msg = str(exc_info.value)
|
|
assert "could not connect to PostgreSQL" in msg
|
|
# Credentials must not appear in the surfaced message
|
|
assert "secret" not in msg
|
|
# Only the first line of the OperationalError should be present (no DETAIL)
|
|
assert "DETAIL" not in msg
|
|
|
|
|
|
def test_pg_introspect_import_error():
|
|
"""If psycopg is missing, introspect_postgres raises ImportError."""
|
|
with patch.dict("sys.modules", {"psycopg": None}):
|
|
with pytest.raises(ImportError, match="psycopg is required") as exc_info:
|
|
introspect_postgres("postgresql://localhost/db")
|
|
# #1906: the PyPI package is graphifyy (double-y), so the install hint must match
|
|
assert "graphifyy[postgres]" in str(exc_info.value)
|
|
|
|
|
|
def test_pg_introspect_grammar_error_surfaces():
|
|
"""If extract_sql reports an error (e.g. tree-sitter-sql missing), the
|
|
introspection must raise instead of returning an empty graph — otherwise
|
|
the CLI prints "PostgreSQL: 0 nodes" as if the database were empty."""
|
|
mock_psycopg = _make_mock_psycopg(
|
|
[("public", "users", "BASE TABLE")], [], [], []
|
|
)
|
|
with patch.dict("sys.modules", {"psycopg": mock_psycopg}):
|
|
with patch(
|
|
"graphify.pg_introspect.extract_sql",
|
|
return_value={
|
|
"nodes": [],
|
|
"edges": [],
|
|
"error": "tree_sitter_sql not installed. Run: pip install tree-sitter-sql",
|
|
},
|
|
):
|
|
with pytest.raises(ImportError, match="tree_sitter_sql not installed"):
|
|
introspect_postgres("postgresql://myhost/mydb")
|
|
|
|
|
|
def test_pg_introspect_uri_forward_slashes():
|
|
"""Assert that the virtual path in postgresql introspection output uses forward slashes on all platforms."""
|
|
mock_psycopg = _make_mock_psycopg([], [], [], [], host="some-host", dbname="some-db")
|
|
with patch.dict("sys.modules", {"psycopg": mock_psycopg}):
|
|
res = introspect_postgres("postgresql://some-host/some-db")
|
|
|
|
# We should have at least the file node
|
|
assert len(res["nodes"]) > 0
|
|
for node in res["nodes"]:
|
|
assert "\\" not in node["source_file"]
|
|
assert "postgresql:/some-host/some-db" in node["source_file"]
|
|
|