chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,331 @@
|
||||
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_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"):
|
||||
introspect_postgres("postgresql://localhost/db")
|
||||
|
||||
|
||||
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"]
|
||||
|
||||
Reference in New Issue
Block a user