fix(sql): recover CREATE POLICY via regex fallback; add pg_policy introspection (#3401)

graphify/extractors/sql.py

@@ -74,6 +74,51 @@
 # Closing it would need a real string heuristic (e.g. an end-of-line opening
 # quote), judged not worth the swallow risk in a recovery-only path.
 
+# Recovers CREATE POLICY statements. The grammar has NO rule for CREATE
+# POLICY at all — not a partial/error-node case like routines, every
+# policy statement disintegrates into loose top-level tokens plus an
+# ERROR node with no CREATE text in it (#3401). So there is no walk-time
+# node to dispatch on; this is whole-file-fallback only, same gate and
+# masking as _ROUTINE_RECOVERY_RX.
+#
+# TO/USING/WITH CHECK are all optional in real SQL (a bare
+# `CREATE POLICY p ON t;` is valid, if useless), so each clause is its
+# own non-greedy optional group rather than a single big alternation —
+# a required-TO assumption silently dropped USING-only policies in an
+
+_POLICY_RECOVERY_RX = re.compile(
+    r"\bCREATE\s+POLICY\s+"
+    r"(\"(?:[^\"\n]|\"\")+\"|[\w$]+)\s+ON\s+"
+    r"((?:\"(?:[^\"\n]|\"\")+\"|[\w$]+)(?:\s*\.\s*(?:\"(?:[^\"\n]|\"\")+\"|[\w$]+))*)"
+    r"(?:\s+AS\s+(PERMISSIVE|RESTRICTIVE))?"
+    r"(?:\s+FOR\s+(SELECT|INSERT|UPDATE|DELETE|ALL))?"
+    r"(?:\s+TO\s+((?:(?:\"(?:[^\"\n]|\"\")+\"|[\w$]+)\s*,\s*)*(?:\"(?:[^\"\n]|\"\")+\"|[\w$]+)))?",
+    re.IGNORECASE,
+)
+# USING (...) / WITH CHECK (...) bodies are located separately below via a
+# manual balanced-paren scan, not captured in this regex -- arbitrarily
+# nested parens (e.g. fn(a, (b + (c))) style expressions) defeat any
+# fixed-depth pattern that only tolerates one level of nesting.
+_USING_KW_RX = re.compile(r"\s*USING\s*\(", re.IGNORECASE)
+_CHECK_KW_RX = re.compile(r"\s*WITH\s+CHECK\s*\(", re.IGNORECASE)
+
+_SQL_PREDICATE_KEYWORDS = {
+    "in", "not", "and", "or", "is", "exists", "any", "all", "some",
+    "case", "when", "between", "like", "ilike", "similar",
+}
+
+def _match_balanced_parens(s, open_pos):
+    depth = 0
+    for i in range(open_pos, len(s)):
+        if s[i] == "(":
+            depth += 1
+        elif s[i] == ")":
+            depth -= 1
+            if depth == 0:
+                return s[open_pos + 1 : i], i + 1
+    return None
+
+_FUNC_CALL_RX = re.compile(r"\b([\w$]+(?:\.[\w$]+)?)\s*\(")
 
 def _scan_sql(text: str) -> tuple[str, list[tuple[int, int]]]:
     """Blank comment and string-literal spans, preserving every offset.
@@ -689,6 +734,42 @@ def _collect_defined_names(node) -> None:
                 continue
             fn_name = m.group(1)
             fn_line = src_text[: m.start()].count("\n") + 1
-            _add_node(_make_id(stem, fn_name), f"{fn_name}()", fn_line)
+            fn_nid = _make_id(stem, fn_name)
+            _add_node(fn_nid, f"{fn_name}()", fn_line)
+            table_nids.setdefault(_norm_ident(fn_name), fn_nid)
+
+        for m in _POLICY_RECOVERY_RX.finditer(masked_src):
+            if any(s <= m.start() < e for s, e in ident_spans):
+                continue
+            pol_name = m.group(1).strip('"')
+            tbl_name = m.group(2)
+            pol_line = src_text[: m.start()].count("\n") + 1
+            tbl_nid = table_nids.get(_norm_ident(tbl_name)) or _ref_stub(tbl_name)
+            pol_nid = _make_id(stem, f"{tbl_name}.{pol_name}")
+            _add_node(pol_nid, pol_name, pol_line)
+            _add_edge(pol_nid, tbl_nid, "applies_to", pol_line)
+            pos = m.end()
+            using_body = check_body = ""
+            um = _USING_KW_RX.match(masked_src, pos)
+            if um:
+                result = _match_balanced_parens(masked_src, um.end() - 1)
+                if result:
+                    using_body, pos = result
+                else:
+                    pos = um.end()
+            cm = _CHECK_KW_RX.match(masked_src, pos)
+            if cm:
+                result = _match_balanced_parens(masked_src, cm.end() - 1)
+                if result:
+                    check_body, pos = result
+            body = " ".join(filter(None, [using_body, check_body]))
+            seen_fns = set()
+            for fm in _FUNC_CALL_RX.finditer(body):
+                fn_key = _norm_ident(fm.group(1))
+                if fn_key in _SQL_PREDICATE_KEYWORDS or fn_key in seen_fns:
+                    continue
+                seen_fns.add(fn_key)
+                fn_nid = table_nids.get(fn_key) or _ref_stub(fm.group(1))
+                _add_edge(pol_nid, fn_nid, "references", pol_line)
 
     return {"nodes": nodes, "edges": edges}

graphify/pg_introspect.py

@@ -96,6 +96,32 @@ def introspect_postgres(dsn: str | None = None) -> dict:
                 ORDER BY ns.nspname, rel.relname, con.conname;
             """)
             fks = cur.fetchall()
+
+            cur.execute("""
+                SELECT
+                    pol.polname AS name,
+                    ns.nspname AS schema,
+                    rel.relname AS table_name,
+                    pol.polcmd AS command,
+                    pol.polpermissive AS permissive,
+                    COALESCE(
+                        (SELECT ARRAY_AGG(
+                                    CASE WHEN ro.oid = 0 THEN 'public' ELSE r.rolname END
+                                    ORDER BY CASE WHEN ro.oid = 0 THEN 'public' ELSE r.rolname END
+                                 )
+                           FROM UNNEST(pol.polroles) AS ro(oid)
+                           LEFT JOIN pg_catalog.pg_roles r ON r.oid = ro.oid),
+                        ARRAY['public']
+                    ) AS roles,
+                    pg_get_expr(pol.polqual, pol.polrelid) AS using_expr,
+                    pg_get_expr(pol.polwithcheck, pol.polrelid) AS check_expr
+                FROM pg_catalog.pg_policy pol
+                JOIN pg_catalog.pg_class rel ON rel.oid = pol.polrelid
+                JOIN pg_catalog.pg_namespace ns ON ns.oid = rel.relnamespace
+                WHERE ns.nspname NOT IN ('pg_catalog', 'information_schema')
+                ORDER BY ns.nspname, rel.relname, pol.polname;
+            """)
+            policies = cur.fetchall()
     finally:
         conn.close()
 
@@ -144,6 +170,20 @@ def introspect_postgres(dsn: str | None = None) -> dict:
                 f" AS $gfx$ {actual_body} $gfx$ LANGUAGE {lang};"
             )
 
+    _CMD_MAP = {"r": "SELECT", "a": "INSERT", "w": "UPDATE", "d": "DELETE", "*": "ALL"}
+    for name, schema, table, cmd, permissive, roles, using_expr, check_expr in policies:
+        clauses = [
+            f"CREATE POLICY {_quote_ident(name)} ON {_quote_ident(schema)}.{_quote_ident(table)}",
+            "AS PERMISSIVE" if permissive else "AS RESTRICTIVE",
+            f"FOR {_CMD_MAP.get(cmd, 'ALL')}",
+            f"TO {', '.join('PUBLIC' if r.lower() == 'public' else _quote_ident(r) for r in roles)}",
+        ]
+        if using_expr:
+            clauses.append(f"USING ({using_expr})")
+        if check_expr:
+            clauses.append(f"WITH CHECK ({check_expr})")
+        ddl.append(" ".join(clauses) + ";")
+
     ddl_string = "\n".join(ddl)
 
     # Determine host/dbname for virtual path DSN sanitization

tests/fixtures/policies.sql

@@ -0,0 +1,12 @@
+create table public.employees (id uuid primary key, home_branch_id uuid);
+
+create or replace function app.is_admin() returns boolean
+  language sql stable as $$ select true $$;
+
+create policy employees_select on public.employees
+  for select to authenticated
+  using (app.is_admin());
+
+create policy employees_update on public.employees
+  for update to authenticated
+  using (app.is_admin()) with check (app.is_admin());

tests/test_create_policy_snippet.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+def test_create_policy_recovered_with_references(tmp_path):
+    """CREATE POLICY has no grammar rule (#3401): every policy statement
+    disintegrates into loose tokens + an ERROR node with no CREATE text in
+    it, so this is whole-file-regex-fallback only, not a walk-time recovery.
+    Asserts both policies land as nodes, each with an applies_to edge to
+    its table and a references edge to the function used in USING/CHECK.
+    """
+    from graphify.extractors.sql import extract_sql
+
+    fixture = Path(__file__).parent / "fixtures" / "policies.sql"
+    result = extract_sql(fixture)
+
+    assert not result.get("error")
+
+    node_labels = {n["label"] for n in result["nodes"]}
+    assert "employees_select" in node_labels
+    assert "employees_update" in node_labels
+
+    def _label(nid):
+        return next(n["label"] for n in result["nodes"] if n["id"] == nid)
+
+    applies_to_targets = {
+        _label(e["target"])
+        for e in result["edges"]
+        if e["relation"] == "applies_to"
+    }
+    assert "public.employees" in applies_to_targets
+
+    references_targets = {
+        _label(e["target"])
+        for e in result["edges"]
+        if e["relation"] == "references"
+        and _label(e["source"]) in node_labels & {"employees_select", "employees_update"}
+    }
+    assert any("is_admin" in t for t in references_targets)