diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 5672bf0df..7151da092 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -16,6 +16,7 @@ Each stage is a single function in its own module. They communicate through plai |--------|----------|----------------| | `detect.py` | `collect_files(root)` | directory → `[Path]` filtered list | | `extract.py` | `extract(path)` | file path → `{nodes, edges}` dict | +| `pg_introspect.py` | `introspect_postgres(dsn)` | PostgreSQL DSN → DDL → `extract_sql()` → `{nodes, edges}` dict | | `build.py` | `build_graph(extractions)` | list of extraction dicts → `nx.Graph` | | `cluster.py` | `cluster(G)` | graph → graph with `community` attr on each node | | `analyze.py` | `analyze(G)` | graph → analysis dict (god nodes, surprises, questions) | diff --git a/README.md b/README.md index 1e4467e74..452bbc237 100644 --- a/README.md +++ b/README.md @@ -166,6 +166,7 @@ Install only what you need: | `gemini` | Google Gemini API | `uv tool install "graphifyy[gemini]"` | | `bedrock` | AWS Bedrock (uses IAM, no API key) | `uv tool install "graphifyy[bedrock]"` | | `sql` | SQL schema extraction | `uv tool install "graphifyy[sql]"` | +| `postgres` | Live PostgreSQL schema introspection | `uv tool install "graphifyy[postgres]"` | | `dm` | BYOND DreamMaker `.dm`/`.dme` AST extraction (may need a C compiler + `python3-dev` if no wheel matches your platform) | `uv tool install "graphifyy[dm]"` | | `chinese` | Chinese query segmentation (jieba) | `uv tool install "graphifyy[chinese]"` | | `all` | Everything above | `uv tool install "graphifyy[all]"` | @@ -523,6 +524,8 @@ graphify extract ./docs --no-cluster # raw extraction only, skip clust graphify extract ./docs --force # overwrite graph.json even if new graph has fewer nodes (use after refactors or to clear ghost duplicates) graphify extract ./docs --dedup-llm # LLM tiebreaker for ambiguous entity pairs (uses same API key) graphify extract ./docs --global --as myrepo # extract and register into the cross-project global graph +graphify extract --postgres "postgresql://user:pass@host:5432/mydb" # introspect a live Postgres database (no LLM key needed) +graphify extract ./src --postgres "postgresql://..." # combine code + live DB schema in one graph GRAPHIFY_MAX_OUTPUT_TOKENS=32768 graphify extract ./docs --backend claude # raise output cap for dense corpora graphify export callflow-html # graphify-out/-callflow.html diff --git a/graphify/__main__.py b/graphify/__main__.py index 8ac1f9f85..68b630a75 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -1797,6 +1797,9 @@ def main() -> None: print(" --out DIR output dir (default: ); writes /graphify-out/") print(" --google-workspace export .gdoc/.gsheet/.gslides shortcuts via gws before extraction") print(" --no-cluster skip clustering, write raw extraction only") + print(" --postgres DSN extract schema from a live PostgreSQL database") + print(" maps tables, views, functions + FK relationships;") + print(" column-level detail is not represented in the graph") print(" --global also merge the resulting graph into the global graph") print(" --as repo tag for --global (default: target directory name)") print(" global add add/update a project graph in the global graph (~/.graphify/global-graph.json)") @@ -3417,20 +3420,26 @@ def _load_graph(p: str): "Usage: graphify extract [--backend gemini|kimi|claude|openai|deepseek|ollama] " "[--model M] [--mode deep] [--out DIR] [--google-workspace] [--no-cluster] " "[--max-workers N] [--token-budget N] [--max-concurrency N] " - "[--api-timeout S]", + "[--api-timeout S] [--postgres DSN]", file=sys.stderr, ) sys.exit(1) - target = Path(sys.argv[2]).resolve() - if not target.exists(): - print(f"error: path not found: {target}", file=sys.stderr) - sys.exit(1) + has_path = True + if sys.argv[2].startswith("-"): + has_path = False + target = Path(".").resolve() + else: + target = Path(sys.argv[2]).resolve() + if not target.exists(): + print(f"error: path not found: {target}", file=sys.stderr) + sys.exit(1) backend: str | None = None model: str | None = None extract_mode: str | None = None out_dir: Path | None = None + cli_postgres_dsn: str | None = None no_cluster = False dedup_llm = False google_workspace = False @@ -3468,7 +3477,7 @@ def _parse_float(name: str, raw: str) -> float: sys.exit(2) return v - args = sys.argv[3:] + args = sys.argv[3:] if has_path else sys.argv[2:] i = 0 while i < len(args): a = args[i] @@ -3526,9 +3535,17 @@ def _parse_float(name: str, raw: str) -> float: cli_excludes.append(args[i + 1]); i += 2 elif a.startswith("--exclude="): cli_excludes.append(a.split("=", 1)[1]); i += 1 + elif a == "--postgres" and i + 1 < len(args): + cli_postgres_dsn = args[i + 1]; i += 2 + elif a.startswith("--postgres="): + cli_postgres_dsn = a.split("=", 1)[1]; i += 1 else: i += 1 + if not has_path and cli_postgres_dsn is None: + print("error: must specify a path to scan or a --postgres DSN", file=sys.stderr) + sys.exit(1) + _VALID_MODES = {"deep"} if extract_mode is not None and extract_mode not in _VALID_MODES: print( @@ -3562,7 +3579,7 @@ def _parse_float(name: str, raw: str) -> float: ) if backend is None: backend = _detect_backend() - if backend is None: + if backend is None and cli_postgres_dsn is None: print( "error: no LLM API key found. Set GEMINI_API_KEY or GOOGLE_API_KEY " "(gemini), MOONSHOT_API_KEY (kimi), ANTHROPIC_API_KEY (claude), " @@ -3571,14 +3588,14 @@ def _parse_float(name: str, raw: str) -> float: file=sys.stderr, ) sys.exit(1) - if backend not in _BACKENDS: + if backend is not None and backend not in _BACKENDS: print( f"error: unknown backend '{backend}'. " f"Available: {', '.join(sorted(_BACKENDS))}", file=sys.stderr, ) sys.exit(1) - if not _get_backend_api_key(backend): + if backend is not None and not _get_backend_api_key(backend): # Ollama on a loopback URL ignores auth entirely; don't block # the run just because OLLAMA_API_KEY is unset (issue #792). # extract_files_direct already prints a warning and substitutes @@ -3636,9 +3653,17 @@ def _parse_float(name: str, raw: str) -> float: ) manifest_path = graphify_out / "manifest.json" existing_graph_path = graphify_out / "graph.json" - incremental_mode = manifest_path.exists() and existing_graph_path.exists() + incremental_mode = manifest_path.exists() and existing_graph_path.exists() if has_path else False - if incremental_mode: + if not has_path: + code_files = [] + doc_files = [] + paper_files = [] + image_files = [] + deleted_files = [] + unchanged_total = 0 + files_by_type = {} + elif incremental_mode: print(f"[graphify extract] incremental scan of {target}") detection = _detect_incremental( target, @@ -3646,12 +3671,7 @@ def _parse_float(name: str, raw: str) -> float: google_workspace=google_workspace or None, extra_excludes=cli_excludes or None, ) - else: - print(f"[graphify extract] scanning {target}") - detection = _detect(target, google_workspace=google_workspace or None, extra_excludes=cli_excludes or None) - - files_by_type = detection.get("files", {}) - if incremental_mode: + files_by_type = detection.get("files", {}) new_by_type = detection.get("new_files", {}) code_files = [Path(p) for p in new_by_type.get("code", [])] doc_files = [Path(p) for p in new_by_type.get("document", [])] @@ -3660,6 +3680,9 @@ def _parse_float(name: str, raw: str) -> float: deleted_files = list(detection.get("deleted_files", [])) unchanged_total = sum(len(v) for v in detection.get("unchanged_files", {}).values()) else: + print(f"[graphify extract] scanning {target}") + detection = _detect(target, google_workspace=google_workspace or None, extra_excludes=cli_excludes or None) + files_by_type = detection.get("files", {}) code_files = [Path(p) for p in files_by_type.get("code", [])] doc_files = [Path(p) for p in files_by_type.get("document", [])] paper_files = [Path(p) for p in files_by_type.get("paper", [])] @@ -3790,13 +3813,25 @@ def _progress(idx: int, total: int, _result: dict) -> None: sem_result["input_tokens"] += fresh.get("input_tokens", 0) sem_result["output_tokens"] += fresh.get("output_tokens", 0) - # Merge AST + semantic. Order matters for deduplication: passing AST + pg_result: dict = {"nodes": [], "edges": []} + if cli_postgres_dsn is not None: + from graphify.pg_introspect import introspect_postgres + print(f"[graphify extract] introspecting PostgreSQL schema...") + try: + pg_result = introspect_postgres(cli_postgres_dsn) + except (ConnectionError, ImportError) as exc: + print(f"error: {exc}", file=sys.stderr) + sys.exit(1) + print(f"[graphify extract] PostgreSQL: {len(pg_result['nodes'])} nodes, " + f"{len(pg_result['edges'])} edges") + + # Merge AST + semantic + pg_result. Order matters for deduplication: passing AST # first means semantic node attributes win on collision (richer labels # for symbols also referenced in docs). Hyperedges only come from the # semantic side. merged: dict = { - "nodes": list(ast_result.get("nodes", [])) + list(sem_result.get("nodes", [])), - "edges": list(ast_result.get("edges", [])) + list(sem_result.get("edges", [])), + "nodes": list(ast_result.get("nodes", [])) + list(sem_result.get("nodes", [])) + list(pg_result.get("nodes", [])), + "edges": list(ast_result.get("edges", [])) + list(sem_result.get("edges", [])) + list(pg_result.get("edges", [])), "hyperedges": list(sem_result.get("hyperedges", [])), "input_tokens": ast_result.get("input_tokens", 0) + sem_result.get("input_tokens", 0), "output_tokens": ast_result.get("output_tokens", 0) + sem_result.get("output_tokens", 0), diff --git a/graphify/extract.py b/graphify/extract.py index 7b392d6c5..a08dd44b4 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -4716,7 +4716,7 @@ def walk(node, module_nid: str | None = None) -> None: return {"nodes": nodes, "edges": edges} -def extract_sql(path: Path) -> dict: +def extract_sql(path: Path, content: str | bytes | None = None) -> dict: """Extract tables, views, functions, and relationships from .sql files via tree-sitter.""" try: import tree_sitter_sql as tssql @@ -4727,12 +4727,17 @@ def extract_sql(path: Path) -> dict: try: language = Language(tssql.language()) parser = Parser(language) - source = path.read_bytes() + source = ( + content.encode("utf-8") if isinstance(content, str) + else content if content is not None + else path.read_bytes() + ) tree = parser.parse(source) root = tree.root_node except Exception as e: return {"nodes": [], "edges": [], "error": str(e)} + stem = _file_stem(path) str_path = str(path) file_nid = _make_id(str_path) diff --git a/graphify/pg_introspect.py b/graphify/pg_introspect.py new file mode 100644 index 000000000..9182ffeae --- /dev/null +++ b/graphify/pg_introspect.py @@ -0,0 +1,142 @@ +from __future__ import annotations +from pathlib import Path +from graphify.extract import extract_sql + + +def _quote_ident(name: str) -> str: + """Double-quote a PostgreSQL identifier, escaping embedded double-quotes.""" + return '"' + name.replace('"', '""') + '"' + + +def introspect_postgres(dsn: str | None = None) -> dict: + """Connect to PostgreSQL, reconstruct DDL, and extract via extract_sql().""" + try: + import psycopg + except ModuleNotFoundError: + raise ImportError( + "psycopg is required for --postgres. " + "Install with: pip install 'graphify[postgres]'" + ) + + try: + conn = psycopg.connect(dsn or "") # empty string = PG* env vars + except psycopg.OperationalError as exc: + # Sanitize: strip the DSN/credentials that psycopg may embed in the + # OperationalError message (e.g. "connection to server … failed: …\nDETAIL: …") + msg = str(exc).split("\n")[0] + raise ConnectionError(f"could not connect to PostgreSQL: {msg}") from None + + try: + conn.execute("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE READ ONLY DEFERRABLE") + + # 1. Query tables + with conn.cursor() as cur: + cur.execute(""" + SELECT table_schema, table_name, table_type + FROM information_schema.tables + WHERE table_schema NOT IN ('pg_catalog', 'information_schema') + ORDER BY table_schema, table_name; + """) + tables = cur.fetchall() + + # 2. Query views + cur.execute(""" + SELECT table_schema, table_name, view_definition + FROM information_schema.views + WHERE table_schema NOT IN ('pg_catalog', 'information_schema') + ORDER BY table_schema, table_name; + """) + views = cur.fetchall() + + # 3. Query routines (functions/procedures), including language + cur.execute(""" + SELECT routine_schema, routine_name, routine_type, + routine_definition, external_language + FROM information_schema.routines + WHERE routine_schema NOT IN ('pg_catalog', 'information_schema') + ORDER BY routine_schema, routine_name; + """) + routines = cur.fetchall() + + # 4. Query foreign keys — grouped by constraint to handle composites + cur.execute(""" + SELECT + tc.constraint_name, + kcu1.table_schema, + kcu1.table_name, + ARRAY_AGG(kcu1.column_name ORDER BY kcu1.ordinal_position) AS columns, + kcu2.table_schema AS foreign_table_schema, + kcu2.table_name AS foreign_table_name, + ARRAY_AGG(kcu2.column_name ORDER BY kcu2.ordinal_position) AS foreign_columns + FROM + information_schema.table_constraints AS tc + JOIN information_schema.referential_constraints AS rc + ON tc.constraint_name = rc.constraint_name + AND tc.table_schema = rc.constraint_schema + JOIN information_schema.key_column_usage AS kcu1 + ON tc.constraint_name = kcu1.constraint_name + AND tc.table_schema = kcu1.table_schema + JOIN information_schema.key_column_usage AS kcu2 + ON rc.unique_constraint_name = kcu2.constraint_name + AND rc.unique_constraint_schema = kcu2.table_schema + AND kcu1.position_in_unique_constraint = kcu2.ordinal_position + WHERE tc.constraint_type = 'FOREIGN KEY' + AND tc.table_schema NOT IN ('pg_catalog', 'information_schema') + GROUP BY tc.constraint_name, kcu1.table_schema, kcu1.table_name, + kcu2.table_schema, kcu2.table_name + ORDER BY kcu1.table_schema, kcu1.table_name; + """) + fks = cur.fetchall() + finally: + conn.close() + + ddl = [] + + # Tables — quote identifiers to handle reserved words, hyphens, mixed-case + for schema, name, ttype in tables: + if ttype == "BASE TABLE": + ddl.append(f"CREATE TABLE {_quote_ident(schema)}.{_quote_ident(name)} (id INT);") + + # Views — real body if available, stub if NULL (permission denied) + for schema, name, body in views: + if body: + ddl.append(f"CREATE VIEW {_quote_ident(schema)}.{_quote_ident(name)} AS {body};") + else: + ddl.append(f"CREATE VIEW {_quote_ident(schema)}.{_quote_ident(name)} AS SELECT 1;") + + # 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. + for schema, name, rtype, body, ext_lang in routines: + lang = (ext_lang or "plpgsql").lower() + fn_sig = f"{_quote_ident(schema)}.{_quote_ident(name)}()" + stub_body = "BEGIN SELECT 1; END;" + if rtype in ("FUNCTION", "PROCEDURE"): + actual_body = body if body else stub_body + # Represent PROCEDUREs as FUNCTION so tree-sitter-sql can parse them + ddl.append( + f"CREATE FUNCTION {fn_sig} RETURNS void" + 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 + info = psycopg.conninfo.conninfo_to_dict(dsn or "") + host = info.get("host", "localhost") + dbname = info.get("dbname", "db") + virtual_path = Path(f"postgresql://{host}/{dbname}") + + # Pass virtual path and in-memory DDL content to extract_sql + result = extract_sql(virtual_path, content=ddl_string) + return result \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 8c4c071d5..2c7597fdb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ svg = ["matplotlib"] leiden = ["graspologic; python_version < '3.13'"] office = ["python-docx", "openpyxl"] google = ["openpyxl"] +postgres = ["psycopg[binary]"] video = ["faster-whisper; python_version >= '3.11'", "yt-dlp"] kimi = ["openai", "tiktoken"] ollama = ["openai"] diff --git a/tests/test_pg_introspect.py b/tests/test_pg_introspect.py new file mode 100644 index 000000000..a1c5a29c8 --- /dev/null +++ b/tests/test_pg_introspect.py @@ -0,0 +1,275 @@ +import sys +from unittest.mock import MagicMock, patch +import pytest +from pathlib import Path +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(). + """ + + 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 + + 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 "information_schema.referential_constraints" 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, + } + 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 "()" + 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_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")