From b78c44992ac4429d1aadecd60f3edcb2a563bc93 Mon Sep 17 00:00:00 2001 From: Tanay <2021cd25ta@mitsgwl.ac.in> Date: Mon, 1 Jun 2026 20:40:00 +0530 Subject: [PATCH 1/2] feat: Add live PostgreSQL introspection connector (resolves #1093) --- ARCHITECTURE.md | 1 + README.md | 3 + graphify/__main__.py | 69 ++++++++++++++------ graphify/extract.py | 9 ++- graphify/pg_introspect.py | 122 ++++++++++++++++++++++++++++++++++++ pyproject.toml | 1 + tests/test_pg_introspect.py | 105 +++++++++++++++++++++++++++++++ 7 files changed, 288 insertions(+), 22 deletions(-) create mode 100644 graphify/pg_introspect.py create mode 100644 tests/test_pg_introspect.py 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..7e8d8cf1f 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -1797,6 +1797,7 @@ 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 PostgreSQL connection string (live schema introspection)") 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 +3418,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 +3475,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 +3533,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 +3577,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 +3586,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 +3651,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 +3669,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 +3678,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 +3811,21 @@ 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...") + pg_result = introspect_postgres(cli_postgres_dsn) + 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..d6adee60d --- /dev/null +++ b/graphify/pg_introspect.py @@ -0,0 +1,122 @@ +from __future__ import annotations +from pathlib import Path +from graphify.extract import extract_sql + +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]'" + ) + + conn = psycopg.connect(dsn or "") # empty string = PG* env vars + 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) + cur.execute(""" + SELECT routine_schema, routine_name, routine_type, routine_definition + 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 + cur.execute(""" + SELECT + kcu1.table_schema, + kcu1.table_name, + kcu1.column_name, + kcu2.table_schema AS foreign_table_schema, + kcu2.table_name AS foreign_table_name, + kcu2.column_name AS foreign_column_name + 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') + ORDER BY kcu1.table_schema, kcu1.table_name, kcu1.ordinal_position; + """) + fks = cur.fetchall() + finally: + conn.close() + + ddl = [] + + # Tables + for schema, name, ttype in tables: + if ttype == "BASE TABLE": + ddl.append(f"CREATE TABLE {schema}.{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 {schema}.{name} AS {body};") + else: + ddl.append(f"CREATE VIEW {schema}.{name} AS SELECT 1;") + + # Functions & Procedures — real body if available, stub if NULL + for schema, name, rtype, body in routines: + if rtype == "FUNCTION": + if body: + ddl.append(f"CREATE FUNCTION {schema}.{name}() RETURNS void AS $$ {body} $$ LANGUAGE plpgsql;") + else: + ddl.append(f"CREATE FUNCTION {schema}.{name}() RETURNS void AS $$ BEGIN SELECT 1; END; $$ LANGUAGE plpgsql;") + elif rtype == "PROCEDURE": + if body: + # To make procedures extractable by tree-sitter-sql (which does not support CREATE PROCEDURE), + # we represent them as CREATE FUNCTION in the reconstructed DDL. + ddl.append(f"CREATE FUNCTION {schema}.{name}() RETURNS void AS $$ {body} $$ LANGUAGE plpgsql;") + else: + ddl.append(f"CREATE FUNCTION {schema}.{name}() RETURNS void AS $$ BEGIN SELECT 1; END; $$ LANGUAGE plpgsql;") + + # FK edges + for t_schema, t_name, col, r_schema, r_name, r_col in fks: + ddl.append( + f"ALTER TABLE {t_schema}.{t_name} " + f"ADD CONSTRAINT fk_{t_schema}_{t_name}_{col} FOREIGN KEY ({col}) REFERENCES {r_schema}.{r_name}({r_col});" + ) + + 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..567fe3a9c --- /dev/null +++ b/tests/test_pg_introspect.py @@ -0,0 +1,105 @@ +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 + +def test_pg_introspect_success(): + # Canned database catalog details + mock_tables = [ + ("public", "users", "BASE TABLE"), + ("public", "orders", "BASE TABLE"), + ] + mock_views = [ + ("public", "active_users", "SELECT * FROM public.users WHERE active = true"), + ] + mock_routines = [ + ("public", "calculate_total", "FUNCTION", "SELECT 42;"), + ("public", "do_nothing", "PROCEDURE", None), + ] + mock_fks = [ + ("public", "orders", "user_id", "public", "users", "id"), + ] + + # Mock psycopg cursors and connections + 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 mock_tables + elif "information_schema.views" in q: + return mock_views + elif "information_schema.routines" in q: + return mock_routines + elif "information_schema.referential_constraints" in q: + return mock_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 = "host=myhost dbname=mydb user=myuser password=mypassword" + return info_mock + + mock_connect = MagicMock(return_value=MockConnection()) + + mock_psycopg = MagicMock() + mock_psycopg.connect = mock_connect + mock_psycopg.conninfo.conninfo_to_dict = MagicMock(return_value={ + "host": "myhost", + "dbname": "mydb", + }) + + with patch.dict("sys.modules", {"psycopg": mock_psycopg}): + res = introspect_postgres("postgresql://myuser:mypassword@myhost/mydb") + + # Assertions + # 1. Check validate_extraction passes + errors = validate_extraction(res) + assert errors == [], f"Validation errors: {errors}" + + # 2. Check source_file is sanitized virtual path (without 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. Check correct node labels were created + node_labels = {n["label"] for n in res["nodes"]} + assert "public.users" in node_labels + assert "public.orders" in node_labels + assert "public.active_users" in node_labels + assert "public.calculate_total()" in node_labels + assert "public.do_nothing()" in node_labels + + # 4. Verify edge relationships (contains edges for all items, and references edge from orders to users) + file_nodes = [n for n in res["nodes"] if n["file_type"] == "code" and n["label"] == "mydb"] + assert len(file_nodes) == 1 + + # Verify orders references users edge exists + users_nid = [n["id"] for n in res["nodes"] if n["label"] == "public.users"][0] + orders_nid = [n["id"] for n in res["nodes"] if n["label"] == "public.orders"][0] + + 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 + +def test_pg_introspect_import_error(): + # If psycopg module is missing or cannot be imported, introspect_postgres raises ImportError + with patch.dict("sys.modules", {"psycopg": None}): + with pytest.raises(ImportError, match="psycopg is required"): + introspect_postgres("postgresql://localhost/db") From d8f79c1c58ad66fd6f63b032e677a5187df7c71d Mon Sep 17 00:00:00 2001 From: Tanay <2021cd25ta@mitsgwl.ac.in> Date: Tue, 2 Jun 2026 11:03:51 +0530 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20address=20PR=20review=20=E2=80=94=20?= =?UTF-8?q?quote=20identifiers,=20composite=20FK=20grouping,=20connection?= =?UTF-8?q?=20error=20handling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Quote all PostgreSQL identifiers with _quote_ident() to handle reserved words, hyphens, mixed-case, and embedded double-quotes - GROUP BY constraint in FK query + ARRAY_AGG to emit one ALTER TABLE per composite FK (prevents duplicate references edges) - Catch psycopg.OperationalError and re-raise as sanitized ConnectionError (no credentials in message, no raw stack trace) - Use $gfx$ dollar-quote tag to avoid collision with $$ in function bodies - Read external_language from information_schema.routines instead of hardcoding plpgsql - Add tests for quoted/special-char table names and composite FKs - Add docs note that column-level detail is not represented in the graph --- graphify/__main__.py | 10 +- graphify/pg_introspect.py | 80 ++++++----- tests/test_pg_introspect.py | 262 +++++++++++++++++++++++++++++------- 3 files changed, 274 insertions(+), 78 deletions(-) diff --git a/graphify/__main__.py b/graphify/__main__.py index 7e8d8cf1f..68b630a75 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -1797,7 +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 PostgreSQL connection string (live schema introspection)") + 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)") @@ -3815,7 +3817,11 @@ def _progress(idx: int, total: int, _result: dict) -> None: if cli_postgres_dsn is not None: from graphify.pg_introspect import introspect_postgres print(f"[graphify extract] introspecting PostgreSQL schema...") - pg_result = introspect_postgres(cli_postgres_dsn) + 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") diff --git a/graphify/pg_introspect.py b/graphify/pg_introspect.py index d6adee60d..9182ffeae 100644 --- a/graphify/pg_introspect.py +++ b/graphify/pg_introspect.py @@ -2,6 +2,12 @@ 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: @@ -12,10 +18,17 @@ def introspect_postgres(dsn: str | None = None) -> dict: "Install with: pip install 'graphify[postgres]'" ) - conn = psycopg.connect(dsn or "") # empty string = PG* env vars + 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(""" @@ -35,24 +48,26 @@ def introspect_postgres(dsn: str | None = None) -> dict: """) views = cur.fetchall() - # 3. Query routines (functions/procedures) + # 3. Query routines (functions/procedures), including language cur.execute(""" - SELECT routine_schema, routine_name, routine_type, routine_definition + 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 + # 4. Query foreign keys — grouped by constraint to handle composites cur.execute(""" SELECT + tc.constraint_name, kcu1.table_schema, kcu1.table_name, - kcu1.column_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, - kcu2.column_name AS foreign_column_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 @@ -67,7 +82,9 @@ def introspect_postgres(dsn: str | None = None) -> dict: 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') - ORDER BY kcu1.table_schema, kcu1.table_name, kcu1.ordinal_position; + 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: @@ -75,38 +92,41 @@ def introspect_postgres(dsn: str | None = None) -> dict: ddl = [] - # Tables + # 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 {schema}.{name} (id INT);") + 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 {schema}.{name} AS {body};") + ddl.append(f"CREATE VIEW {_quote_ident(schema)}.{_quote_ident(name)} AS {body};") else: - ddl.append(f"CREATE VIEW {schema}.{name} AS SELECT 1;") + ddl.append(f"CREATE VIEW {_quote_ident(schema)}.{_quote_ident(name)} AS SELECT 1;") # Functions & Procedures — real body if available, stub if NULL - for schema, name, rtype, body in routines: - if rtype == "FUNCTION": - if body: - ddl.append(f"CREATE FUNCTION {schema}.{name}() RETURNS void AS $$ {body} $$ LANGUAGE plpgsql;") - else: - ddl.append(f"CREATE FUNCTION {schema}.{name}() RETURNS void AS $$ BEGIN SELECT 1; END; $$ LANGUAGE plpgsql;") - elif rtype == "PROCEDURE": - if body: - # To make procedures extractable by tree-sitter-sql (which does not support CREATE PROCEDURE), - # we represent them as CREATE FUNCTION in the reconstructed DDL. - ddl.append(f"CREATE FUNCTION {schema}.{name}() RETURNS void AS $$ {body} $$ LANGUAGE plpgsql;") - else: - ddl.append(f"CREATE FUNCTION {schema}.{name}() RETURNS void AS $$ BEGIN SELECT 1; END; $$ LANGUAGE plpgsql;") - - # FK edges - for t_schema, t_name, col, r_schema, r_name, r_col in fks: + # 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 {t_schema}.{t_name} " - f"ADD CONSTRAINT fk_{t_schema}_{t_name}_{col} FOREIGN KEY ({col}) REFERENCES {r_schema}.{r_name}({r_col});" + 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) diff --git a/tests/test_pg_introspect.py b/tests/test_pg_introspect.py index 567fe3a9c..a1c5a29c8 100644 --- a/tests/test_pg_introspect.py +++ b/tests/test_pg_introspect.py @@ -5,101 +5,271 @@ from graphify.pg_introspect import introspect_postgres from graphify.validate import validate_extraction -def test_pg_introspect_success(): - # Canned database catalog details - mock_tables = [ - ("public", "users", "BASE TABLE"), - ("public", "orders", "BASE TABLE"), - ] - mock_views = [ - ("public", "active_users", "SELECT * FROM public.users WHERE active = true"), - ] - mock_routines = [ - ("public", "calculate_total", "FUNCTION", "SELECT 42;"), - ("public", "do_nothing", "PROCEDURE", None), - ] - mock_fks = [ - ("public", "orders", "user_id", "public", "users", "id"), - ] - # Mock psycopg cursors and connections +# --------------------------------------------------------------------------- +# 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 mock_tables + return tables elif "information_schema.views" in q: - return mock_views + return views elif "information_schema.routines" in q: - return mock_routines + return routines elif "information_schema.referential_constraints" in q: - return mock_fks + 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 = "host=myhost dbname=mydb user=myuser password=mypassword" + info_mock.dsn = f"host={host} dbname={dbname} user=myuser password=secret" return info_mock - mock_connect = MagicMock(return_value=MockConnection()) - mock_psycopg = MagicMock() - mock_psycopg.connect = mock_connect - mock_psycopg.conninfo.conninfo_to_dict = MagicMock(return_value={ - "host": "myhost", - "dbname": "mydb", - }) + 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") - # Assertions - # 1. Check validate_extraction passes + # 1. validate_extraction must pass errors = validate_extraction(res) assert errors == [], f"Validation errors: {errors}" - # 2. Check source_file is sanitized virtual path (without credentials) + # 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. Check correct node labels were created + # 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 "public.users" in node_labels - assert "public.orders" in node_labels - assert "public.active_users" in node_labels - assert "public.calculate_total()" in node_labels - assert "public.do_nothing()" in node_labels + 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. Verify edge relationships (contains edges for all items, and references edge from orders to users) + # 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 - # Verify orders references users edge exists - users_nid = [n["id"] for n in res["nodes"] if n["label"] == "public.users"][0] - orders_nid = [n["id"] for n in res["nodes"] if n["label"] == "public.orders"][0] - - 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 + # 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 module is missing or cannot be imported, introspect_postgres raises ImportError + """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")