diff --git a/.gitignore b/.gitignore index 893c0f0..5b743e4 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,9 @@ seed/warehouse.duckdb seed/warehouse.duckdb.wal seed/orders.csv .logs/*.jsonl +.logs/*.sqlite +.logs/*.lock +.logs/approval.key +.logs/demo_audit.jsonl +build/ +dist/ diff --git a/CHANGELOG.md b/CHANGELOG.md index a8a1a9b..7b51681 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,30 @@ All notable changes to **sql-write-gate** are documented here. +## [1.1.0] — 2026-09-06 + +### SQLGuard depth (evolve-in-place) + +- **AST patterns** (`ast_patterns` + `ast_guard`): cartesian / missing-predicate JOINs BLOCK; tautology WHERE (`1=1` / `TRUE`) clears `has_where` so destructive still owns full-table UPDATE/DELETE; DDL dangerous flags on findings. +- **Pluggable rule registry** (`registry.GuardRegistry`): register/unregister guards without forking `engine.py`; built-ins load via `default_registry()`. +- **Permissions** in `policy.yaml` (`permissions.tables` / `allow_tables` + `enforced`): table×operation allowlists (GameStream-style). Offline-safe when unset. +- **Numeric risk** on `Decision`: `risk_score` 0–100 + `risk_factors` (from guard results + AST flags + optional EXPLAIN). +- **Optional EXPLAIN cost** (`explain.py` + `explain_cost` guard): adapters when `conn` present; skips cleanly offline / when threshold unset. +- **Schema hallucination**: unknown table/column → `schema_hallucination` (catalog allowlist); SELECT path checks referenced tables/columns; `allow_unknown_tables/columns` knobs. +- **Audit**: `actor`, `model_id`, `prompt_summary`, `sql`, `risk`/`risk_score`, `latency_ms`, `success`, `decision`. +- **DataPilot BLOCK/EXECUTE**: stable `datapilot` verb on MCP (`datapilot_block_or_execute`), CLI `datapilot` / `serve` (HTTP `POST /v1/check|/v1/execute`), Python `write_gate.datapilot` / `write_gate.api`. +- Preserve freshness + PII behavior from 1.0.x. + +### Docs / tests + +- README: testable/demoable first (AST / registry / permissions / risk / DataPilot). +- Tests: `tests/test_v110.py`; v1.0.1 suite accepts `>= 1.0.1`. + +### Notes + +- Keep **非生产唯一边界 / 非唯一边界**. +- Package version **1.1.0**. + ## [1.0.1] — 2026-09-06 ### Fixes (pilot-ready / 非唯一边界 unchanged) diff --git a/README.md b/README.md index 98840d3..d83948f 100644 --- a/README.md +++ b/README.md @@ -1,21 +1,21 @@ -# sql-write-gate +# sql-write-gate (SQLGuard) [![CI](https://github.com/tangyf07/sql-write-gate/actions/workflows/ci.yml/badge.svg)](https://github.com/tangyf07/sql-write-gate/actions/workflows/ci.yml) [![Release](https://img.shields.io/github/v/release/tangyf07/sql-write-gate)](https://github.com/tangyf07/sql-write-gate/releases/latest) [![Python 3.11+](https://img.shields.io/badge/python-3.11%2B-blue.svg)](https://www.python.org/downloads/) [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE) -写库前门禁 · Policy firewall for AI agents writing to databases. +写库前门禁 · **SQLGuard** — policy firewall for AI agents writing to databases. Prevent Claude Code, Codex, Cursor and MCP agents from executing unsafe database operations. ``` - Agent SQL ──► sql-write-gate ──► ALLOW / BLOCK / APPROVAL ──► Database + DataPilot / Agent SQL ──► SQLGuard (sql-write-gate) ──► ALLOW / BLOCK / APPROVAL ──► Database ``` Deterministic policy engine (sqlglot AST + catalog + policy.yaml). **No LLM. No API key.** -> **v1.0.1 — pilot-ready on the declared support matrix** (DuckDB / PostgreSQL / MySQL / SQLite + listed SQL features + entrypoints below). +> **v1.1.0 — SQLGuard** — stronger AST analysis, permissions, risk scores, schema hallucination block, DataPilot HTTP API; pilot-ready on the declared support matrix (DuckDB / PostgreSQL / MySQL / SQLite + listed SQL features + entrypoints below). > **非生产唯一边界 / 非唯一边界** — **not** the sole production DB security boundary. Combine with least-privilege DB roles, network isolation, and human workflows. > **未列语法拒绝** — unsupported / ambiguous SQL → REJECT/BLOCK (`unsupported_sql`, fail closed), never silent ALLOW as read-only. @@ -41,6 +41,18 @@ sql-write-gate check "DELETE FROM users" # → BLOCKED rule=delete_without_where ``` +## Try it + +```bash +make install && make test && make demo +sql-write-gate check "DELETE FROM orders" +# → BLOCKED delete_without_where +sql-write-gate datapilot --json "SELECT o.order_id FROM orders o CROSS JOIN orders p" +# → {"datapilot": "BLOCK", "rule_id": "cartesian_join", ...} +sql-write-gate serve --port 8787 +# POST /v1/check {"sql": "..."} → action + risk_score + datapilot +``` + ## Entrypoints (stable) | Entrypoint | Role | @@ -51,6 +63,7 @@ sql-write-gate check "DELETE FROM users" | CLI `proxy` | Gate then execute if ALLOW | | CLI `approve` / `resolve` / `reject` | Human approve / recover (trusted executor + token) | | CLI `audit` / `pending` / `init` / `exec` | Ops helpers | +| CLI `serve` | SQLGuard DataPilot HTTP (`/v1/check`, `/v1/execute`) | ```bash sql-write-gate check "SQL" # evaluate SQL; no execute @@ -63,6 +76,39 @@ sql-write-gate audit # TIME / SOURCE / OP / TABLE / VERDICT sql-write-gate init # scaffold policy.yaml + catalog.json ``` +## DataPilot API contract (SQLGuard) + +DataPilot calls this gate **outbound**. Prefer MCP `query_sql` / `write_sql`, CLI `check` / `exec` / `proxy`, or HTTP: + +```bash +sql-write-gate serve --host 127.0.0.1 --port 8787 +``` + +| Method | Path | Behavior | +|--------|------|----------| +| `GET` | `/healthz` | `{ok, product: SQLGuard, version}` | +| `POST` | `/v1/check` | Evaluate only → `action` + `risk_score` (`executed: false`) | +| `POST` | `/v1/execute` | Gate then execute **only on ALLOW** | +| `POST` | `/v1/block` | Alias of `/v1/check` | + +Request JSON: `{ "sql": "...", "actor"?, "model_id"?, "prompt_summary"?, "database"?, "policy"? }`. + +Response always includes `action` (`ALLOW` \| `BLOCK` \| `REQUIRE_APPROVAL`), `rule_id`, `reason`, `risk_score`, `risk_factors`, `executed`. Treat anything other than `ALLOW` as non-executing. + +### GameStream-style permissions + +```yaml +permissions: + enforce: true + tables: + orders: [select, insert, update] +hallucination: + allow_unknown_tables: false + allow_unknown_columns: false +``` + +Python alias: `from write_gate import sqlguard` (product helpers); package/CLI names unchanged. + ## Declared databases (supported) | Backend | How to connect | Notes | @@ -91,7 +137,11 @@ Anything **not** in this matrix (other warehouses, wire-protocol proxies, distri ## What it does (on the matrix) - `DROP` / `TRUNCATE` / `ALTER` → BLOCK -- `DELETE` / `UPDATE` without `WHERE` → BLOCK +- `DELETE` / `UPDATE` without `WHERE` (incl. tautology `WHERE 1=1` / `TRUE`) → BLOCK +- Cartesian / missing-predicate JOINs → BLOCK (`cartesian_join`) +- Unknown tables/columns → BLOCK (`schema_hallucination`) with evidence +- GameStream-style `permissions.tables` allowlists in `policy.yaml` +- Numeric `risk_score` (0–100) + `risk_factors` on every Decision - Blast-radius COUNT vs `update_rows` / `delete_rows` (dialect quoting; fail-closed on estimate error) - Schema / PII / restricted columns; PII `SELECT` → REQUIRE_APPROVAL (approve executes once) - Freshness partitions (`dt`); range / NOT / OR / UPSERT SET expired → BLOCK diff --git a/examples/policy.demo.yaml b/examples/policy.demo.yaml index 22ee8eb..7e39893 100644 --- a/examples/policy.demo.yaml +++ b/examples/policy.demo.yaml @@ -10,3 +10,11 @@ rules: limits: update_rows: 10000 delete_rows: 10000 +# SQLGuard: GameStream-style table allowlist (demo allows orders fully) +permissions: + enforce: true + tables: + orders: [select, insert, update, delete] +hallucination: + allow_unknown_tables: false + allow_unknown_columns: false diff --git a/examples/policy.yaml b/examples/policy.yaml index 197fe22..0ac097a 100644 --- a/examples/policy.yaml +++ b/examples/policy.yaml @@ -10,3 +10,11 @@ rules: limits: update_rows: 100 delete_rows: 50 +# SQLGuard 1.1 — GameStream-style table permissions (DataPilot / GameStream) +permissions: + enforce: true + tables: + orders: [select, insert, update] +hallucination: + allow_unknown_tables: false + allow_unknown_columns: false diff --git a/policy.yaml b/policy.yaml index 197fe22..9d17421 100644 --- a/policy.yaml +++ b/policy.yaml @@ -10,3 +10,10 @@ rules: limits: update_rows: 100 delete_rows: 50 +permissions: + enforce: true + tables: + orders: [select, insert, update] +hallucination: + allow_unknown_tables: false + allow_unknown_columns: false diff --git a/pyproject.toml b/pyproject.toml index 9e61ed1..eea371b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "sql-write-gate" -version = "1.0.1" -description = "Policy firewall for AI agents writing to databases (sql-write-gate)" +version = "1.1.0" +description = "SQLGuard — policy firewall for AI agents writing to databases (sql-write-gate)" readme = "README.md" license = { text = "MIT" } requires-python = ">=3.11" diff --git a/scripts/demo.py b/scripts/demo.py index 54897b4..77fdb5c 100644 --- a/scripts/demo.py +++ b/scripts/demo.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Print the three canonical gate cases: legal / expired / PII. No API key.""" +"""SQLGuard demo: classic three cases + AST / hallucination / risk / audit.""" from __future__ import annotations @@ -16,11 +16,28 @@ from write_gate.paths import DB_PATH, DEMO_POLICY_PATH # noqa: E402 from write_gate.wrapper import WriteGate # noqa: E402 - CASES = [ ("用例 1 · 合法写入（新鲜分区 + 非 PII 列）", LEGAL_WRITE_SQL, True, "ok"), ("用例 2 · 过期分区写入", EXPIRED_WRITE_SQL, False, "expired_partition"), ("用例 3 · PII 列写入", PII_WRITE_SQL, False, "pii_column"), + ( + "用例 4 · AST · DELETE WHERE 1=1（全表写）", + "DELETE FROM orders WHERE 1=1", + False, + "delete_without_where", + ), + ( + "用例 5 · AST · CROSS JOIN（笛卡尔积）", + "SELECT o.order_id FROM orders o CROSS JOIN orders p", + False, + "cartesian_join", + ), + ( + "用例 6 · Schema hallucination · 未知表", + "SELECT * FROM game_stream_sessions", + False, + "schema_hallucination", + ), ] @@ -33,7 +50,18 @@ def _banner(title: str) -> None: def main() -> int: rc = 0 - with WriteGate(db_path=DB_PATH, policy_path=DEMO_POLICY_PATH) as gate: + audit_path = ROOT / ".logs" / "demo_audit.jsonl" + if audit_path.exists(): + audit_path.unlink() + with WriteGate( + db_path=DB_PATH, + policy_path=DEMO_POLICY_PATH, + audit_path=audit_path, + agent="demo", + actor="sqlguard-demo", + model_id="demo-offline", + prompt_summary="make demo SQLGuard cases", + ) as gate: for title, sql, expect_allow, expect_rule in CASES: _banner(title) print(f"SQL:\n {sql}") @@ -41,6 +69,8 @@ def main() -> int: verdict = "ALLOWED" if evidence.allowed else "BLOCKED" print(f"VERDICT: {verdict}") print(f"rule_id: {evidence.rule_id}") + print(f"risk_score: {evidence.risk_score}") + print(f"risk_factors: {evidence.risk_factors}") print(f"message: {evidence.message}") print("evidence:") print(json.dumps(evidence.to_dict(), ensure_ascii=False, indent=2)) @@ -57,8 +87,28 @@ def main() -> int: ) rc = 1 print() + + _banner("Audit · last lines (actor / risk_score / latency_ms)") + if audit_path.is_file(): + lines = audit_path.read_text(encoding="utf-8").strip().splitlines() + for line in lines[-4:]: + rec = json.loads(line) + print( + json.dumps( + { + "actor": rec.get("actor"), + "decision": rec.get("decision"), + "rule_id": rec.get("rule_id"), + "risk_score": rec.get("risk_score"), + "latency_ms": rec.get("latency_ms"), + "success": rec.get("success"), + "sql": (rec.get("sql") or "")[:60], + }, + ensure_ascii=False, + ) + ) if rc == 0: - print("demo: three cases matched expected verdicts") + print("demo: SQLGuard cases matched expected verdicts") return rc diff --git a/scripts/demo_walkthrough.py b/scripts/demo_walkthrough.py index ccbb04a..8d01425 100644 --- a/scripts/demo_walkthrough.py +++ b/scripts/demo_walkthrough.py @@ -115,8 +115,26 @@ def _count_order(db_path: Path, order_id: int) -> int: conn.close() +def _ensure_demo_approval_trust() -> None: + """Trusted-executor token for approve/resolve in this offline demo.""" + key = ROOT / ".logs" / "approval.key" + key.parent.mkdir(parents=True, exist_ok=True) + secret = "sqlguard-demo-approval-secret" + if not key.is_file(): + key.write_text(secret + "\n", encoding="utf-8") + try: + key.chmod(0o600) + except OSError: + pass + else: + secret = key.read_text(encoding="utf-8").strip() or secret + os.environ["SQL_WRITE_GATE_APPROVAL_KEY_FILE"] = str(key) + os.environ["SQL_WRITE_GATE_APPROVAL_TOKEN"] = secret + + def main() -> int: os.chdir(ROOT) + _ensure_demo_approval_trust() gate = _gate_cmd() print(f"walkthrough CLI: {gate[0]}") diff --git a/src/write_gate/__init__.py b/src/write_gate/__init__.py index 8e37f69..1b4f7cc 100644 --- a/src/write_gate/__init__.py +++ b/src/write_gate/__init__.py @@ -1,7 +1,7 @@ -"""sql-write-gate: policy firewall for AI agents writing to databases.""" +"""sql-write-gate (SQLGuard): policy firewall for AI agents writing to databases.""" from write_gate.decision import Decision, Evidence from write_gate.wrapper import WriteGate -__all__ = ["WriteGate", "Evidence", "Decision"] -__version__ = "1.0.1" +__all__ = ["WriteGate", "Evidence", "Decision", "__version__"] +__version__ = "1.1.0" diff --git a/src/write_gate/api.py b/src/write_gate/api.py new file mode 100644 index 0000000..d45ca87 --- /dev/null +++ b/src/write_gate/api.py @@ -0,0 +1,185 @@ +"""DataPilot HTTP API: clear BLOCK / EXECUTE surface (stdlib http.server). + +Contract (JSON):: + + POST /v1/check {"sql": "...", "actor"?, "model_id"?, "prompt_summary"?} + POST /v1/execute same body — gate then execute on ALLOW only + GET /healthz {"ok": true, "product": "SQLGuard", "version": "..."} + +Responses always include ``action`` (ALLOW|BLOCK|REQUIRE_APPROVAL), +``risk_score``, ``risk_factors``, ``rule_id``, ``reason``, and for execute +``executed`` (bool). DataPilot should treat anything other than ALLOW as +non-executing; REQUIRE_APPROVAL may include ``approval_id``. +""" + +from __future__ import annotations + +import json +import sys +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any +from urllib.parse import urlparse + +from write_gate import __version__ +from write_gate.decision import ACTION_ALLOW +from write_gate.wrapper import WriteGate + + +def decision_response( + decision: Any, + *, + executed: bool = False, + rows: list | None = None, + rowcount: int | None = None, +) -> dict[str, Any]: + from write_gate.datapilot import to_datapilot_action + + payload = decision.to_dict() + payload["executed"] = bool(executed) + payload["product"] = "SQLGuard" + payload["datapilot"] = to_datapilot_action(decision) # BLOCK | EXECUTE | APPROVAL + if rows is not None: + payload["rows"] = rows + if rowcount is not None: + payload["rowcount"] = rowcount + return payload + + +def _gate_from_body(body: dict[str, Any], *, defaults: dict[str, Any]) -> WriteGate: + return WriteGate( + database=body.get("database") or defaults.get("database"), + db_path=body.get("db_path") or defaults.get("db_path"), + catalog_path=body.get("catalog") or defaults.get("catalog"), + policy_path=body.get("policy") or defaults.get("policy"), + agent=body.get("agent") or defaults.get("agent") or "datapilot", + actor=body.get("actor") or body.get("agent") or "datapilot", + model_id=body.get("model_id"), + prompt_summary=body.get("prompt_summary"), + ) + + +def handle_datapilot_request( + method: str, + path: str, + body: dict[str, Any] | None, + *, + defaults: dict[str, Any] | None = None, +) -> tuple[int, dict[str, Any]]: + """Pure request handler used by the HTTP server and unit tests.""" + defaults = defaults or {} + parsed = urlparse(path) + route = parsed.path.rstrip("/") or "/" + + if method == "GET" and route in {"/healthz", "/v1/healthz"}: + return 200, { + "ok": True, + "product": "SQLGuard", + "version": __version__, + "package": "sql-write-gate", + } + + if method != "POST" or route not in {"/v1/check", "/v1/execute", "/v1/block"}: + return 404, {"error": "not_found", "path": route} + + body = body or {} + sql = body.get("sql") + if not sql or not isinstance(sql, str): + return 400, {"error": "missing_sql", "action": "BLOCK"} + + # /v1/block is an alias that always evaluates (same as check) — DataPilot + # may call it when it only wants a verdict without execute intent. + do_execute = route == "/v1/execute" + with _gate_from_body(body, defaults=defaults) as gate: + if do_execute: + decision, result = gate.execute(sql) + executed = decision.action == ACTION_ALLOW and result is not None + rows = None + rowcount = None + if executed and result is not None: + try: + from write_gate.results import materialize_result + + mat = materialize_result(result) + if mat: + rows = mat.get("rows") + rowcount = mat.get("rowcount") + except Exception: + rowcount = getattr(result, "rowcount", None) + return 200, decision_response( + decision, executed=executed, rows=rows, rowcount=rowcount + ) + decision = gate.check(sql) + return 200, decision_response(decision, executed=False) + + +class DataPilotHandler(BaseHTTPRequestHandler): + server_defaults: dict[str, Any] = {} + + def log_message(self, fmt: str, *args: Any) -> None: # noqa: A003 + sys.stderr.write("sqlguard-api: " + (fmt % args) + "\n") + + def _read_json(self) -> dict[str, Any]: + length = int(self.headers.get("Content-Length") or 0) + if length <= 0: + return {} + raw = self.rfile.read(length) + try: + data = json.loads(raw.decode("utf-8")) + except json.JSONDecodeError: + return {} + return data if isinstance(data, dict) else {} + + def _send(self, status: int, payload: dict[str, Any]) -> None: + body = json.dumps(payload, ensure_ascii=False).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self) -> None: # noqa: N802 + status, payload = handle_datapilot_request( + "GET", self.path, None, defaults=self.server_defaults + ) + self._send(status, payload) + + def do_POST(self) -> None: # noqa: N802 + status, payload = handle_datapilot_request( + "POST", + self.path, + self._read_json(), + defaults=self.server_defaults, + ) + self._send(status, payload) + + +def serve( + host: str = "127.0.0.1", + port: int = 8787, + *, + defaults: dict[str, Any] | None = None, +) -> ThreadingHTTPServer: + """Start DataPilot HTTP server (blocking via serve_forever in caller).""" + DataPilotHandler.server_defaults = defaults or {} + httpd = ThreadingHTTPServer((host, port), DataPilotHandler) + return httpd + + +def run_serve_cli( + host: str = "127.0.0.1", + port: int = 8787, + *, + defaults: dict[str, Any] | None = None, +) -> int: + httpd = serve(host, port, defaults=defaults) + sys.stderr.write( + f"SQLGuard DataPilot API listening on http://{host}:{port} " + f"(POST /v1/check|/v1/execute)\n" + ) + try: + httpd.serve_forever() + except KeyboardInterrupt: + sys.stderr.write("\nshutting down\n") + finally: + httpd.server_close() + return 0 diff --git a/src/write_gate/ast_patterns.py b/src/write_gate/ast_patterns.py new file mode 100644 index 0000000..b5a910c --- /dev/null +++ b/src/write_gate/ast_patterns.py @@ -0,0 +1,193 @@ +"""AST pattern helpers: joins, tautology WHERE, referenced tables/columns. + +Builds on sqlglot Expression trees from ``parser.parse`` — never regex. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from sqlglot import exp + +from write_gate.idents import ident, literal_value + + +@dataclass +class JoinInfo: + """One JOIN (including implicit comma joins).""" + + kind: str # CROSS | INNER | LEFT | RIGHT | FULL | COMMA | UNKNOWN + has_predicate: bool + is_cartesian: bool + right_table: str | None = None + sql: str = "" + + +@dataclass +class AstFindings: + """Structured AST signals surfaced on ParsedSQL / guards.""" + + tables: list[str] = field(default_factory=list) + columns: list[str] = field(default_factory=list) + joins: list[JoinInfo] = field(default_factory=list) + cartesian_joins: list[JoinInfo] = field(default_factory=list) + missing_where: bool = False + tautology_where: bool = False + full_table_write: bool = False + is_ddl: bool = False + is_dml: bool = False + dangerous_flags: list[str] = field(default_factory=list) + + def to_evidence(self) -> dict[str, Any]: + return { + "tables": list(self.tables), + "columns": list(self.columns), + "joins": [ + { + "kind": j.kind, + "has_predicate": j.has_predicate, + "is_cartesian": j.is_cartesian, + "right_table": j.right_table, + } + for j in self.joins + ], + "cartesian_joins": len(self.cartesian_joins), + "missing_where": self.missing_where, + "tautology_where": self.tautology_where, + "full_table_write": self.full_table_write, + "is_ddl": self.is_ddl, + "is_dml": self.is_dml, + "dangerous_flags": list(self.dangerous_flags), + } + + +def _join_kind(join: exp.Join) -> str: + kind = join.args.get("kind") + if kind is not None: + text = str(kind).upper() + if text: + return text + # Implicit comma join: JOIN with no kind and no ON/USING often from FROM a, b + if join.args.get("on") is None and join.args.get("using") is None: + # sqlglot represents comma joins as Join with kind=None + side = join.args.get("side") + if side is None and kind is None: + return "COMMA" + return "INNER" + + +def analyze_joins(stmt: exp.Expression) -> list[JoinInfo]: + out: list[JoinInfo] = [] + for join in stmt.find_all(exp.Join): + on = join.args.get("on") + using = join.args.get("using") + has_pred = on is not None or using is not None + kind = _join_kind(join) + is_cross = kind == "CROSS" or (kind == "COMMA" and not has_pred) + is_cartesian = is_cross or (not has_pred and kind in {"COMMA", "INNER", "UNKNOWN"}) + # Explicit INNER/LEFT with ON is fine + if has_pred: + is_cartesian = False + if kind == "CROSS": + is_cartesian = True + right = ident(join.this) if join.this is not None else None + out.append( + JoinInfo( + kind=kind, + has_predicate=has_pred, + is_cartesian=is_cartesian, + right_table=right, + sql=join.sql()[:200], + ) + ) + return out + + +def is_tautology_predicate(node: exp.Expression | None) -> bool: + """True for WHERE TRUE / 1 / 1=1 / '1'='1' style always-true filters.""" + if node is None: + return False + if isinstance(node, exp.Paren): + return is_tautology_predicate(node.this) + if isinstance(node, exp.Boolean): + return bool(node.this) is True + if isinstance(node, exp.Literal): + if node.is_int: + try: + return int(node.this) != 0 + except (TypeError, ValueError): + return False + # Non-empty string literal alone is unusual; treat "true" case-insensitively + return str(node.this).lower() in {"true", "t", "yes"} + if isinstance(node, exp.EQ): + left = literal_value(node.this) + right = literal_value(node.expression) + if left is None or right is None: + return False + return left == right and not isinstance(node.this, exp.Column) and not isinstance( + node.expression, exp.Column + ) + return False + + +def referenced_tables(stmt: exp.Expression) -> list[str]: + names: list[str] = [] + for table in stmt.find_all(exp.Table): + name = ident(table) + if name and name not in names: + names.append(name) + return names + + +def referenced_columns(stmt: exp.Expression) -> list[str]: + names: list[str] = [] + for col in stmt.find_all(exp.Column): + name = ident(col) + if name and name != "*" and name not in names: + names.append(name) + return names + + +def analyze_statement( + stmt: exp.Expression | None, + *, + operation: str, + has_where: bool, + where: exp.Expression | None, +) -> AstFindings: + findings = AstFindings() + if stmt is None: + return findings + + findings.tables = referenced_tables(stmt) + findings.columns = referenced_columns(stmt) + findings.joins = analyze_joins(stmt) + findings.cartesian_joins = [j for j in findings.joins if j.is_cartesian] + findings.is_ddl = operation == "ddl" + findings.is_dml = operation in {"insert", "update", "delete"} + + tautology = is_tautology_predicate(where) + findings.tautology_where = tautology + findings.missing_where = operation in {"update", "delete"} and ( + not has_where or tautology + ) + findings.full_table_write = findings.missing_where + + flags: list[str] = [] + if findings.is_ddl: + flags.append("ddl") + if findings.full_table_write: + flags.append("full_table_write") + if findings.tautology_where and operation in {"update", "delete"}: + flags.append("tautology_where") + if findings.cartesian_joins: + flags.append("cartesian_join") + if isinstance(stmt, (exp.Drop, exp.TruncateTable)) or type(stmt).__name__ in { + "TruncateTable", + "Alter", + "AlterTable", + }: + flags.append("destructive_ddl") + findings.dangerous_flags = flags + return findings diff --git a/src/write_gate/audit.py b/src/write_gate/audit.py index b3ffc4a..6d16c9b 100644 --- a/src/write_gate/audit.py +++ b/src/write_gate/audit.py @@ -184,6 +184,11 @@ def append_audit( execution_outcome: str | None = None, error_class: str | None = None, request_id: str | None = None, + actor: str | None = None, + model_id: str | None = None, + prompt_summary: str | None = None, + latency_ms: float | None = None, + success: bool | None = None, ) -> None: """Append one audit JSONL record with correlatable ids (v0.23). @@ -200,10 +205,18 @@ def append_audit( from write_gate.runtime import new_request_id rid = new_request_id(request_id) + # success: explicit override, else infer from executed / action + if success is None: + if executed is not None: + success = bool(executed) + else: + success = decision.action == "ALLOW" + record: dict[str, Any] = { "timestamp": datetime.now(timezone.utc).isoformat(), "request_id": rid, "agent": agent, + "actor": actor or agent, "environment": environment, "sql": decision.sql, "operation": decision.operation, @@ -211,6 +224,10 @@ def append_audit( "estimated_rows": decision.estimated_rows, "decision": decision.action, "rule_id": decision.rule_id, + "risk": decision.risk, + "risk_score": getattr(decision, "risk_score", 0), + "risk_factors": list(getattr(decision, "risk_factors", None) or []), + "success": success, } if decision.approval_id: record["approval_id"] = decision.approval_id @@ -222,6 +239,12 @@ def append_audit( record["execution_outcome"] = execution_outcome if error_class is not None: record["error_class"] = error_class + if model_id is not None: + record["model_id"] = model_id + if prompt_summary is not None: + record["prompt_summary"] = str(prompt_summary)[:500] + if latency_ms is not None: + record["latency_ms"] = round(float(latency_ms), 3) dest = Path(path) if path else default_audit_path() dest.parent.mkdir(parents=True, exist_ok=True) maybe_rotate(dest) @@ -273,19 +296,25 @@ def format_audit_table(rows: Iterable[dict[str, Any]]) -> str: records = list(rows) if not records: return EMPTY_AUDIT_MESSAGE - headers = ("TIME", "SOURCE", "OP", "TABLE", "VERDICT", "RULE") + show_risk = any(rec.get("risk_score") is not None for rec in records) + if show_risk: + headers = ("TIME", "SOURCE", "OP", "TABLE", "VERDICT", "RULE", "RISK") + else: + headers = ("TIME", "SOURCE", "OP", "TABLE", "VERDICT", "RULE") extracted: list[tuple[str, ...]] = [] for rec in records: - extracted.append( - ( - format_audit_time(rec.get("timestamp")), - str(rec.get("agent") or "-"), - str(rec.get("operation") or "-"), - str(rec.get("table") or "-"), - format_verdict(rec.get("decision")), - str(rec.get("rule_id") or "-"), - ) + base = ( + format_audit_time(rec.get("timestamp")), + str(rec.get("actor") or rec.get("agent") or "-"), + str(rec.get("operation") or "-"), + str(rec.get("table") or "-"), + format_verdict(rec.get("decision")), + str(rec.get("rule_id") or "-"), ) + if show_risk: + score = rec.get("risk_score") + base = base + (str(score) if score is not None else "-",) + extracted.append(base) widths = [len(h) for h in headers] for row in extracted: for i, cell in enumerate(row): diff --git a/src/write_gate/cli.py b/src/write_gate/cli.py index 28d2ca2..20e0a55 100644 --- a/src/write_gate/cli.py +++ b/src/write_gate/cli.py @@ -44,11 +44,15 @@ def format_decision(decision: Decision) -> str: lines = [ _headline(decision.action), f"Risk: {_safe(decision.risk)}", + f"Risk score: {getattr(decision, 'risk_score', 0)}", f"Operation: {_safe(decision.operation).upper()}", f"Table: {_safe(decision.table)}", f"Rule: {_safe(decision.rule_id)}", f"Reason: {_safe(decision.reason)}", ] + factors = getattr(decision, "risk_factors", None) or [] + if factors: + lines.append(f"Risk factors: {', '.join(factors)}") if decision.estimated_rows is not None: lines.append(f"Estimated rows: {decision.estimated_rows}") if decision.approval_id: @@ -79,13 +83,17 @@ def _require_trust_or_exit() -> int | None: def _gate_from_args(args: argparse.Namespace) -> WriteGate: + agent = getattr(args, "agent", None) or "cli" return WriteGate( db_path=Path(args.db) if getattr(args, "db", None) else None, database=getattr(args, "database", None), catalog_path=Path(args.catalog) if getattr(args, "catalog", None) else None, policy_path=Path(args.policy) if getattr(args, "policy", None) else None, approvals_path=_approvals_path(args), - agent=getattr(args, "agent", None) or "cli", + agent=agent, + actor=getattr(args, "actor", None) or agent, + model_id=getattr(args, "model_id", None), + prompt_summary=getattr(args, "prompt_summary", None), ) @@ -180,6 +188,14 @@ def _add_shared(parser: argparse.ArgumentParser) -> None: ), ) parser.add_argument("--agent", default="cli", help="Audit agent name") + parser.add_argument("--actor", default=None, help="Audit actor (DataPilot / human id)") + parser.add_argument("--model-id", dest="model_id", default=None, help="Model id for audit") + parser.add_argument( + "--prompt-summary", + dest="prompt_summary", + default=None, + help="Short prompt summary for audit (truncated)", + ) parser.add_argument("--json", action="store_true", help="Print machine-readable JSON") parser.add_argument( "--approvals", @@ -318,6 +334,26 @@ def build_parser() -> argparse.ArgumentParser: parents=[queue], ) + serve_p = sub.add_parser( + "serve", + help="SQLGuard DataPilot HTTP API (POST /v1/check|/v1/execute)", + parents=[shared], + ) + serve_p.add_argument("--host", default="127.0.0.1", help="Bind host") + serve_p.add_argument("--port", type=int, default=8787, help="Bind port") + + dp = sub.add_parser( + "datapilot", + help="DataPilot BLOCK/EXECUTE evaluate (optional --execute)", + parents=[shared], + ) + dp.add_argument("sql", help="One SQL statement") + dp.add_argument( + "--execute", + action="store_true", + help="Execute on ALLOW (default: check-only)", + ) + init_p = sub.add_parser( "init", help="Scaffold policy.yaml, catalog.json, GETTING_STARTED.md", @@ -575,6 +611,58 @@ def main(argv: list[str] | None = None) -> int: if args.command == "init": return _cmd_init(args) + if args.command == "datapilot": + from write_gate.datapilot import block_or_execute + + payload = block_or_execute( + args.sql, + execute=bool(args.execute), + database=getattr(args, "database", None), + db_path=getattr(args, "db", None), + catalog_path=getattr(args, "catalog", None), + policy_path=getattr(args, "policy", None), + agent=getattr(args, "agent", None) or "datapilot", + actor=getattr(args, "actor", None), + model_id=getattr(args, "model_id", None), + prompt_summary=getattr(args, "prompt_summary", None), + ) + if args.json: + json.dump(payload, sys.stdout, ensure_ascii=False, indent=2) + sys.stdout.write("\n") + else: + lines = [ + str(payload["datapilot"]), + f"Action: {payload['action']}", + f"Rule: {payload.get('rule_id')}", + f"Risk score: {payload.get('risk_score')}", + f"Reason: {payload.get('reason')}", + f"executed: {'yes' if payload.get('executed') else 'no'}", + ] + sys.stdout.write("\n".join(lines) + "\n") + if payload["datapilot"] == "EXECUTE" and ( + not args.execute or payload.get("executed") + ): + return 0 + if payload["datapilot"] == "APPROVAL": + return 1 + return 2 + + if args.command == "serve": + from write_gate.api import run_serve_cli + + defaults = { + "database": getattr(args, "database", None), + "db_path": getattr(args, "db", None), + "catalog": getattr(args, "catalog", None), + "policy": getattr(args, "policy", None), + "agent": getattr(args, "agent", None) or "datapilot", + } + return run_serve_cli( + host=args.host, + port=args.port, + defaults=defaults, + ) + with _gate_from_args(args) as gate: if args.command == "check": decision = gate.check(args.sql) diff --git a/src/write_gate/config.py b/src/write_gate/config.py index 1af06b3..99242c5 100644 --- a/src/write_gate/config.py +++ b/src/write_gate/config.py @@ -56,6 +56,14 @@ class Policy: result_byte_limit: int | None = None audit_max_bytes: int | None = None audit_rotate_daily: bool | None = None + # SQLGuard 1.1 — GameStream-style table permissions + hallucination knobs + table_permissions: dict[str, list[str]] = field(default_factory=dict) + permissions_enforced: bool = False + allow_unknown_tables: bool = False + allow_unknown_columns: bool = False + default_table_ops: list[str] | None = None + explain_cost_threshold: float | None = None + enable_explain: bool = False def rule_for(self, operation: str) -> str: op = (operation or "ddl").lower() @@ -84,6 +92,13 @@ def with_env_approvals_cleared(self) -> "Policy": result_byte_limit=self.result_byte_limit, audit_max_bytes=self.audit_max_bytes, audit_rotate_daily=self.audit_rotate_daily, + table_permissions=dict(self.table_permissions), + permissions_enforced=self.permissions_enforced, + allow_unknown_tables=self.allow_unknown_tables, + allow_unknown_columns=self.allow_unknown_columns, + default_table_ops=(None if self.default_table_ops is None else list(self.default_table_ops)), + explain_cost_threshold=self.explain_cost_threshold, + enable_explain=self.enable_explain, ) @@ -106,6 +121,12 @@ def policy_from_dict(raw: dict[str, Any] | None = None) -> Policy: result_bytes = limits.get("result_bytes", limits.get("result_byte_limit")) audit_max = limits.get("audit_max_bytes", data.get("audit_max_bytes")) audit_daily = limits.get("audit_rotate_daily", data.get("audit_rotate_daily")) + table_perms, enforced = _parse_permissions(data) + hallu = data.get("hallucination") or data.get("schema") or {} + if not isinstance(hallu, dict): + hallu = {} + explain_thr = limits.get("explain_cost_threshold", data.get("explain_cost_threshold")) + enable_explain = data.get("enable_explain", limits.get("enable_explain", False)) return Policy( environment=str(data.get("environment") or PRODUCTION_DEFAULTS["environment"]), rules=_normalize_rules(data.get("rules")), @@ -116,9 +137,43 @@ def policy_from_dict(raw: dict[str, Any] | None = None) -> Policy: result_byte_limit=(None if result_bytes is None else int(result_bytes)), audit_max_bytes=(None if audit_max is None else int(audit_max)), audit_rotate_daily=(None if audit_daily is None else bool(audit_daily)), + table_permissions=table_perms, + permissions_enforced=enforced, + allow_unknown_tables=bool(hallu.get("allow_unknown_tables", data.get("allow_unknown_tables", False))), + allow_unknown_columns=bool(hallu.get("allow_unknown_columns", data.get("allow_unknown_columns", False))), + default_table_ops=([str(x).lower() for x in (data.get("permissions") or {}).get("default_ops")] if isinstance(data.get("permissions"), dict) and (data.get("permissions") or {}).get("default_ops") is not None else None), + explain_cost_threshold=(None if explain_thr is None else float(explain_thr)), + enable_explain=bool(enable_explain), ) + +def _parse_permissions(data: dict[str, Any]) -> tuple[dict[str, list[str]], bool]: + """Parse permissions: tables map and/or allow_tables list.""" + raw = data.get("permissions") or {} + if not isinstance(raw, dict): + return {}, False + table_perms: dict[str, list[str]] = {} + tables = raw.get("tables") or raw.get("table_permissions") or {} + if isinstance(tables, dict): + for k, v in tables.items(): + key = str(k).lower() + if isinstance(v, (list, tuple, set)): + table_perms[key] = [str(x).lower() for x in v] + elif isinstance(v, str): + table_perms[key] = [v.lower()] + elif v is True: + table_perms[key] = ["select", "insert", "update", "delete"] + allow_tables = raw.get("allow_tables") + if isinstance(allow_tables, list): + default_ops = raw.get("default_ops") or ["select", "insert", "update", "delete"] + ops = [str(x).lower() for x in default_ops] + for name in allow_tables: + table_perms.setdefault(str(name).lower(), list(ops)) + enforced = bool(raw.get("enforced", raw.get("enforce", bool(table_perms)))) + return table_perms, enforced + + def load_policy(path: Path | str | None = None) -> Policy: policy_path = Path(path) if path else default_policy_path() if not policy_path.exists(): diff --git a/src/write_gate/datapilot.py b/src/write_gate/datapilot.py new file mode 100644 index 0000000..fc3b1ee --- /dev/null +++ b/src/write_gate/datapilot.py @@ -0,0 +1,204 @@ +"""Stable DataPilot BLOCK / EXECUTE API for agents (MCP / CLI / thin HTTP). + +Maps gate Decisions to a two-verb surface: + - BLOCK — do not run SQL (covers Decision BLOCK) + - EXECUTE — SQL is safe to run / was run (covers Decision ALLOW) + - APPROVAL — human queue (REQUIRE_APPROVAL); not auto-executed + +``check`` never writes. ``execute`` runs only on ALLOW. +""" + +from __future__ import annotations + +import json +import time +from http.server import BaseHTTPRequestHandler, HTTPServer +from typing import Any +from urllib.parse import urlparse + +from write_gate.decision import ACTION_ALLOW, ACTION_APPROVAL, ACTION_BLOCK, Decision +from write_gate.wrapper import WriteGate + +DATAPILOT_BLOCK = "BLOCK" +DATAPILOT_EXECUTE = "EXECUTE" +DATAPILOT_APPROVAL = "APPROVAL" + + +def to_datapilot_action(decision: Decision) -> str: + if decision.action == ACTION_ALLOW: + return DATAPILOT_EXECUTE + if decision.action == ACTION_APPROVAL: + return DATAPILOT_APPROVAL + return DATAPILOT_BLOCK + + +def datapilot_payload( + decision: Decision, + *, + executed: bool = False, + latency_ms: float | None = None, + rows: list[list[Any]] | None = None, + rowcount: int | None = None, + extra: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Stable JSON shape for DataPilot clients.""" + payload: dict[str, Any] = { + "datapilot": to_datapilot_action(decision), + "action": decision.action, # ALLOW / BLOCK / REQUIRE_APPROVAL + "executed": bool(executed), + "rule_id": decision.rule_id, + "reason": decision.reason, + "operation": decision.operation, + "table": decision.table, + "risk": decision.risk, + "risk_score": getattr(decision, "risk_score", 0), + "risk_factors": list(getattr(decision, "risk_factors", []) or []), + "sql": decision.sql, + "evidence": decision.evidence, + "estimated_rows": decision.estimated_rows, + } + if decision.approval_id: + payload["approval_id"] = decision.approval_id + if latency_ms is not None: + payload["latency_ms"] = round(float(latency_ms), 3) + if rowcount is not None: + payload["rowcount"] = rowcount + if rows is not None: + payload["rows"] = rows + if extra: + payload.update(extra) + return payload + + +def block_or_execute( + sql: str, + *, + execute: bool = False, + gate: WriteGate | None = None, + database: str | None = None, + db_path: str | None = None, + catalog_path: str | None = None, + policy_path: str | None = None, + agent: str = "datapilot", + actor: str | None = None, + model_id: str | None = None, + prompt_summary: str | None = None, +) -> dict[str, Any]: + """Evaluate SQL; optionally execute on ALLOW. + + Returns DataPilot payload. Never raises on BLOCK. + """ + owns = gate is None + t0 = time.perf_counter() + g = gate or WriteGate( + database=database, + db_path=db_path, + catalog_path=catalog_path, + policy_path=policy_path, + agent=agent or "datapilot", + actor=actor or agent, + model_id=model_id, + prompt_summary=prompt_summary, + ) + if actor: + g.actor = actor + if model_id is not None: + g.model_id = model_id + if prompt_summary is not None: + g.prompt_summary = prompt_summary + try: + if execute: + decision, result = g.execute(sql) + executed = decision.action == ACTION_ALLOW and result is not None + rows = None + rowcount = None + if executed and result is not None: + try: + from write_gate.results import materialize_result + from write_gate.runtime import load_runtime_settings + + mat = materialize_result(result, settings=load_runtime_settings()) + if mat is not None: + rows = list(mat.get("rows") or []) + rowcount = mat.get("rowcount", len(rows)) + except Exception: + rowcount = getattr(result, "rowcount", None) + latency = (time.perf_counter() - t0) * 1000.0 + return datapilot_payload( + decision, + executed=executed, + latency_ms=latency, + rows=rows, + rowcount=rowcount, + ) + decision = g.check(sql) + latency = (time.perf_counter() - t0) * 1000.0 + return datapilot_payload(decision, executed=False, latency_ms=latency) + finally: + if owns: + close = getattr(g, "close", None) + if callable(close): + close() + + +def serve_http( + host: str = "127.0.0.1", + port: int = 8787, + *, + database: str | None = None, + policy_path: str | None = None, + catalog_path: str | None = None, +) -> None: + """Minimal HTTP DataPilot: POST /v1/datapilot {\"sql\", \"execute\": bool}.""" + + class Handler(BaseHTTPRequestHandler): + def log_message(self, fmt: str, *args) -> None: # quieter + return + + def _json(self, code: int, body: dict[str, Any]) -> None: + raw = json.dumps(body, ensure_ascii=False).encode("utf-8") + self.send_response(code) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(raw))) + self.end_headers() + self.wfile.write(raw) + + def do_GET(self) -> None: # noqa: N802 + path = urlparse(self.path).path + if path in {"/", "/health", "/v1/health"}: + self._json(200, {"ok": True, "service": "sql-write-gate-datapilot"}) + return + self._json(404, {"error": "not_found"}) + + def do_POST(self) -> None: # noqa: N802 + path = urlparse(self.path).path + if path not in {"/v1/datapilot", "/datapilot", "/v1/block-or-execute"}: + self._json(404, {"error": "not_found"}) + return + length = int(self.headers.get("Content-Length") or 0) + raw = self.rfile.read(length) if length else b"{}" + try: + body = json.loads(raw.decode("utf-8") or "{}") + except json.JSONDecodeError: + self._json(400, {"error": "invalid_json"}) + return + sql = str(body.get("sql") or "") + if not sql.strip(): + self._json(400, {"error": "sql_required"}) + return + result = block_or_execute( + sql, + execute=bool(body.get("execute", False)), + database=body.get("database") or database, + policy_path=body.get("policy_path") or policy_path, + catalog_path=body.get("catalog_path") or catalog_path, + actor=body.get("actor"), + model_id=body.get("model_id"), + prompt_summary=body.get("prompt_summary"), + agent=str(body.get("agent") or "datapilot-http"), + ) + self._json(200, result) + + httpd = HTTPServer((host, port), Handler) + print(f"DataPilot HTTP on http://{host}:{port}/v1/datapilot", flush=True) + httpd.serve_forever() diff --git a/src/write_gate/decision.py b/src/write_gate/decision.py index e056864..e5ad71a 100644 --- a/src/write_gate/decision.py +++ b/src/write_gate/decision.py @@ -35,6 +35,12 @@ RULE_ENV = "environment_policy" RULE_RAW_DB_CLI = "raw_db_cli" +RULE_CARTESIAN = "cartesian_join" +RULE_FULL_TABLE = "full_table_write" +RULE_PERMISSION = "table_permission" +RULE_HALLUCINATION = "schema_hallucination" +RULE_EXPLAIN_COST = "explain_cost" + @dataclass class GuardResult: @@ -115,6 +121,8 @@ class Decision: table: str | None = None estimated_rows: int | None = None approval_id: str | None = None + risk_score: int = 0 # 0–100 numeric risk (SQLGuard) + risk_factors: list[str] = field(default_factory=list) @property def allowed(self) -> bool: @@ -138,6 +146,8 @@ def to_dict(self) -> dict[str, Any]: "table": self.table, "estimated_rows": self.estimated_rows, "approval_id": self.approval_id, + "risk_score": self.risk_score, + "risk_factors": list(self.risk_factors), } return payload diff --git a/src/write_gate/engine.py b/src/write_gate/engine.py index 6200d00..a618750 100644 --- a/src/write_gate/engine.py +++ b/src/write_gate/engine.py @@ -2,6 +2,9 @@ any BLOCK → BLOCK; else any APPROVAL → REQUIRE_APPROVAL; else ALLOW. Guard order prefers specific dangerous-SQL rules over environment policy. + +Guards are loaded from ``write_gate.registry.default_registry`` so third-party +rules can register without forking this module. """ from __future__ import annotations @@ -9,6 +12,7 @@ from dataclasses import dataclass, field from typing import Any, Callable +from write_gate.adapters.base import BACKEND_DUCKDB, sqlglot_dialect from write_gate.catalog import Catalog from write_gate.config import Policy, production_policy from write_gate.decision import ( @@ -22,30 +26,12 @@ Decision, GuardResult, ) -from write_gate.guards import ( - check_blast_radius, - check_destructive, - check_environment, - check_freshness, - check_pii, - check_schema, -) -from write_gate.adapters.base import BACKEND_DUCKDB, sqlglot_dialect from write_gate.parser import ParsedSQL, parse +from write_gate.registry import default_registry +from write_gate.risk import score_from_results GuardFn = Callable[["Context"], GuardResult] -# Destructive first so DELETE without WHERE reports delete_without_where -# even when environment also blocks DELETE. -GUARDS: list[GuardFn] = [ - check_destructive, - check_schema, - check_pii, - check_freshness, - check_blast_radius, - check_environment, -] - @dataclass class Context: @@ -67,6 +53,7 @@ def evaluate( dialect: str = BACKEND_DUCKDB, *, human_approved: bool = False, + registry=None, ) -> Decision: parsed = parse(sql, dialect=sqlglot_dialect(dialect)) ctx = Context( @@ -78,11 +65,23 @@ def evaluate( dialect=dialect, human_approved=human_approved, ) - results = [guard(ctx) for guard in GUARDS] + reg = registry if registry is not None else default_registry() + results = reg.run(ctx) ctx.guard_results = results return reduce(ctx, results) +def _guards_list() -> list[GuardFn]: + return default_registry().functions() + + +def __getattr__(name: str): + """Lazy GUARDS list so callers see the live registry.""" + if name == 'GUARDS': + return _guards_list() + raise AttributeError(name) + + def reduce(ctx: Context, results: list[GuardResult]) -> Decision: parsed = ctx.parsed estimated = _first_estimated(results) @@ -96,8 +95,9 @@ def reduce(ctx: Context, results: list[GuardResult]) -> Decision: if blocks: chosen = blocks[0] - return Decision( - action=ACTION_BLOCK, + action = ACTION_BLOCK + decision = Decision( + action=action, risk=chosen.risk, rule_id=chosen.rule_id or RULE_OK, reason=chosen.reason, @@ -107,10 +107,11 @@ def reduce(ctx: Context, results: list[GuardResult]) -> Decision: table=parsed.table, estimated_rows=estimated if estimated is not None else chosen.evidence.get("estimated_rows"), ) - if approvals: + elif approvals: chosen = approvals[0] - return Decision( - action=ACTION_APPROVAL, + action = ACTION_APPROVAL + decision = Decision( + action=action, risk=chosen.risk, rule_id=chosen.rule_id or RULE_OK, reason=chosen.reason, @@ -120,17 +121,32 @@ def reduce(ctx: Context, results: list[GuardResult]) -> Decision: table=parsed.table, estimated_rows=estimated if estimated is not None else chosen.evidence.get("estimated_rows"), ) - return Decision( - action=ACTION_ALLOW, - risk=RISK_LOW, - rule_id=RULE_OK, - reason=_allow_reason(parsed), - evidence=evidence_acc, - sql=ctx.sql, + else: + action = ACTION_ALLOW + decision = Decision( + action=action, + risk=RISK_LOW, + rule_id=RULE_OK, + reason=_allow_reason(parsed), + evidence=evidence_acc, + sql=ctx.sql, + operation=parsed.operation, + table=parsed.table, + estimated_rows=estimated, + ) + + explain_ev = evidence_acc.get("explain_cost") or {} + score, factors = score_from_results( + action=decision.action, operation=parsed.operation, - table=parsed.table, - estimated_rows=estimated, + results=results, + dangerous_flags=getattr(parsed, "dangerous_flags", None), + explain_cost=explain_ev.get("cost"), + explain_cost_threshold=getattr(ctx.policy, "explain_cost_threshold", None), ) + decision.risk_score = score + decision.risk_factors = factors + return decision def _first_estimated(results: list[GuardResult]) -> int | None: diff --git a/src/write_gate/explain.py b/src/write_gate/explain.py new file mode 100644 index 0000000..2165756 --- /dev/null +++ b/src/write_gate/explain.py @@ -0,0 +1,98 @@ +"""Optional EXPLAIN / cost estimation via adapters. Offline-safe degrade.""" + +from __future__ import annotations + +from typing import Any + +from write_gate.adapters.base import BACKEND_DUCKDB, BACKEND_MYSQL, BACKEND_POSTGRES, BACKEND_SQLITE + + +def estimate_cost( + conn: Any | None, + sql: str, + *, + backend: str = BACKEND_DUCKDB, +) -> tuple[float | None, str | None, str | None]: + """Return (cost, plan_text, skip_reason). + + When ``conn`` is missing or EXPLAIN fails, degrade gracefully with a skip + reason — never raise into the gate path. + """ + if conn is None: + return None, None, "no connection" + try: + if backend == BACKEND_POSTGRES: + return _pg_cost(conn, sql) + if backend == BACKEND_MYSQL: + return _mysql_cost(conn, sql) + if backend == BACKEND_SQLITE: + return _sqlite_cost(conn, sql) + return _duckdb_cost(conn, sql) + except Exception as exc: # noqa: BLE001 — degrade only + return None, None, f"explain failed: {exc}" + + +def _duckdb_cost(conn: Any, sql: str) -> tuple[float | None, str | None, str | None]: + # DuckDB: EXPLAIN; no numeric cost — use plan length as soft signal. + cur = conn.execute(f"EXPLAIN {sql}") + rows = cur.fetchall() + plan = "\n".join(str(r[0]) if len(r) == 1 else str(r) for r in rows) + # Heuristic: longer plans / SEQ_SCAN hints → higher soft cost + cost = float(len(plan)) + if "SEQ_SCAN" in plan.upper() or "FULL" in plan.upper(): + cost *= 1.5 + return cost, plan[:2000], None + + +def _pg_cost(conn: Any, sql: str) -> tuple[float | None, str | None, str | None]: + cur = conn.execute(f"EXPLAIN (FORMAT JSON) {sql}") + rows = cur.fetchall() + plan = str(rows[0][0]) if rows else "" + cost = _extract_pg_total_cost(plan) + return cost, plan[:2000], None + + +def _extract_pg_total_cost(plan: str) -> float | None: + import json + import re + + try: + data = json.loads(plan) if plan.lstrip().startswith("[") else None + except json.JSONDecodeError: + data = None + if isinstance(data, list) and data: + node = data[0].get("Plan") if isinstance(data[0], dict) else None + if isinstance(node, dict) and "Total Cost" in node: + return float(node["Total Cost"]) + m = re.search(r"cost=\d+\.\d+\.\.(\d+\.\d+)", plan) + if m: + return float(m.group(1)) + return float(len(plan)) if plan else None + + +def _mysql_cost(conn: Any, sql: str) -> tuple[float | None, str | None, str | None]: + cur = conn.execute(f"EXPLAIN {sql}") + rows = cur.fetchall() + plan = "\n".join(str(r) for r in rows) + # rows column is typically index 8 or named — soft sum + cost = 0.0 + for r in rows: + try: + # PyMySQL dict or tuple + if isinstance(r, dict): + cost += float(r.get("rows") or 0) + elif len(r) > 8 and r[8] is not None: + cost += float(r[8]) + except (TypeError, ValueError): + continue + return (cost if cost else float(len(plan))), plan[:2000], None + + +def _sqlite_cost(conn: Any, sql: str) -> tuple[float | None, str | None, str | None]: + cur = conn.execute(f"EXPLAIN QUERY PLAN {sql}") + rows = cur.fetchall() + plan = "\n".join(str(r) for r in rows) + cost = float(len(rows) * 10 + len(plan)) + if "SCAN" in plan.upper() and "USING" not in plan.upper(): + cost *= 2 + return cost, plan[:2000], None diff --git a/src/write_gate/guards/__init__.py b/src/write_gate/guards/__init__.py index 15435a6..0165432 100644 --- a/src/write_gate/guards/__init__.py +++ b/src/write_gate/guards/__init__.py @@ -1,17 +1,23 @@ """Guard functions. Each returns PASS | WARN | APPROVAL | BLOCK.""" +from write_gate.guards.ast_guard import check_ast_patterns from write_gate.guards.blast_radius import check_blast_radius from write_gate.guards.destructive import check_destructive from write_gate.guards.environment import check_environment +from write_gate.guards.explain_cost import check_explain_cost from write_gate.guards.freshness import check_freshness +from write_gate.guards.permissions import check_permissions from write_gate.guards.pii import check_pii from write_gate.guards.schema import check_schema __all__ = [ + "check_ast_patterns", "check_blast_radius", "check_destructive", "check_environment", + "check_explain_cost", "check_freshness", + "check_permissions", "check_pii", "check_schema", ] diff --git a/src/write_gate/guards/ast_guard.py b/src/write_gate/guards/ast_guard.py new file mode 100644 index 0000000..3028276 --- /dev/null +++ b/src/write_gate/guards/ast_guard.py @@ -0,0 +1,36 @@ +"""AST pattern guard: cartesian / missing-predicate JOINs. + +Missing WHERE / tautology WHERE on UPDATE/DELETE are handled by the +destructive guard after ``parser`` clears ``has_where`` for tautologies. +""" + +from __future__ import annotations + +from write_gate.decision import RULE_CARTESIAN, GuardResult + +NAME = "ast_patterns" + + +def check_ast_patterns(ctx) -> GuardResult: + parsed = ctx.parsed + if parsed.statement is None or parsed.error: + return GuardResult.pass_(NAME) + + findings = getattr(parsed, "findings", None) + evidence = findings.to_evidence() if findings is not None else {} + + if findings is not None and findings.cartesian_joins: + tables = [j.right_table or "?" for j in findings.cartesian_joins] + kinds = sorted({j.kind for j in findings.cartesian_joins}) + return GuardResult.block( + NAME, + RULE_CARTESIAN, + ( + "Dangerous JOIN without predicates " + f"(cartesian / {', '.join(kinds)}) involving {tables}; " + "blocked by AST analysis" + ), + evidence=evidence, + ) + + return GuardResult.pass_(NAME, evidence=evidence) diff --git a/src/write_gate/guards/explain_cost.py b/src/write_gate/guards/explain_cost.py new file mode 100644 index 0000000..3695c43 --- /dev/null +++ b/src/write_gate/guards/explain_cost.py @@ -0,0 +1,49 @@ +"""Optional EXPLAIN cost guard. Skips cleanly when no connection (offline OK).""" + +from __future__ import annotations + +from write_gate.decision import RULE_EXPLAIN_COST, RISK_MEDIUM, GuardResult +from write_gate.explain import estimate_cost + +NAME = "explain_cost" + + +def check_explain_cost(ctx) -> GuardResult: + parsed = ctx.parsed + if parsed.statement is None or parsed.error: + return GuardResult.pass_(NAME) + + enabled = bool(getattr(ctx.policy, "enable_explain", False)) + threshold = getattr(ctx.policy, "explain_cost_threshold", None) + if not enabled and threshold is None: + return GuardResult.pass_( + NAME, + evidence={"skipped": True, "reason": "explain disabled / no threshold"}, + ) + if threshold is None: + # enabled but no threshold → collect evidence only + threshold = float("inf") + + conn = getattr(ctx, "conn", None) + cost, plan, skipped = estimate_cost(conn, ctx.sql, backend=getattr(ctx, "dialect", "duckdb")) + evidence = { + "threshold": threshold, + "cost": cost, + "skipped": skipped, + } + if plan is not None: + evidence["plan_preview"] = plan[:500] + if skipped or cost is None: + return GuardResult.pass_(NAME, evidence=evidence) + if float(cost) >= float(threshold): + return GuardResult.block( + NAME, + RULE_EXPLAIN_COST, + ( + f"EXPLAIN cost {cost} exceeds policy threshold {threshold}; " + "blocked (optional cost guard)" + ), + risk=RISK_MEDIUM, + evidence=evidence, + ) + return GuardResult.pass_(NAME, evidence=evidence) diff --git a/src/write_gate/guards/permissions.py b/src/write_gate/guards/permissions.py new file mode 100644 index 0000000..58a7c25 --- /dev/null +++ b/src/write_gate/guards/permissions.py @@ -0,0 +1,78 @@ +"""Policy-driven table/operation permissions (GameStream-style allowlists).""" + +from __future__ import annotations + +from write_gate.decision import RULE_PERMISSION, GuardResult + +NAME = "permissions" + + +def check_permissions(ctx) -> GuardResult: + parsed = ctx.parsed + if parsed.statement is None or parsed.error: + return GuardResult.pass_(NAME) + + policy = ctx.policy + perms = getattr(policy, "table_permissions", None) or {} + # No permissions map configured → pass (environment / schema still apply). + if not perms and not getattr(policy, "permissions_enforced", False): + return GuardResult.pass_(NAME) + + operation = parsed.operation if parsed.operation != "unknown" else "ddl" + tables = list(getattr(parsed, "tables_referenced", None) or []) + if parsed.table and parsed.table not in tables: + tables.insert(0, parsed.table) + if not tables: + # DDL without parseable table still subject to default deny when enforced + if getattr(policy, "permissions_enforced", False) and operation != "select": + return GuardResult.block( + NAME, + RULE_PERMISSION, + f"{operation.upper()} denied: no target table resolved under permissions policy", + evidence={"operation": operation, "tables": []}, + ) + return GuardResult.pass_(NAME) + + default_ops = getattr(policy, "default_table_ops", None) + for table in tables: + allowed = perms.get(table) + if allowed is None: + if getattr(policy, "allow_unknown_tables", False): + continue + if not perms and default_ops is None: + continue + # Enforced catalog: unknown table relative to permissions map + if getattr(policy, "permissions_enforced", False) or perms: + return GuardResult.block( + NAME, + RULE_PERMISSION, + ( + f"Table {table} is not in the permissions allowlist; " + f"{operation.upper()} denied (GameStream-style policy)" + ), + evidence={ + "table": table, + "operation": operation, + "allowed_tables": sorted(perms.keys()), + }, + ) + continue + ops = {str(o).lower() for o in allowed} + if operation not in ops and "*" not in ops and "all" not in ops: + return GuardResult.block( + NAME, + RULE_PERMISSION, + ( + f"{operation.upper()} on {table} denied by table permissions " + f"(allowed: {sorted(ops)})" + ), + evidence={ + "table": table, + "operation": operation, + "allowed_ops": sorted(ops), + }, + ) + return GuardResult.pass_( + NAME, + evidence={"tables": tables, "operation": operation, "checked": True}, + ) diff --git a/src/write_gate/guards/schema.py b/src/write_gate/guards/schema.py index b581e75..d748199 100644 --- a/src/write_gate/guards/schema.py +++ b/src/write_gate/guards/schema.py @@ -5,7 +5,7 @@ from sqlglot import exp from write_gate.catalog import TableSpec -from write_gate.decision import RULE_SCHEMA, GuardResult +from write_gate.decision import RULE_HALLUCINATION, RULE_SCHEMA, GuardResult from write_gate.parser import literal_value, type_ok NAME = "schema" @@ -27,7 +27,7 @@ def check_schema(ctx) -> GuardResult: operation = parsed.operation if operation == "select": - return GuardResult.pass_(NAME) + return _check_select_hallucination(parsed, ctx) if operation == "ddl": # Destructive/environment guards own DROP/ALTER/TRUNCATE/CREATE. @@ -49,10 +49,10 @@ def check_schema(ctx) -> GuardResult: if spec is None: return GuardResult.block( NAME, - RULE_SCHEMA, - f"未知表 {table_name}，不在目录中", + RULE_HALLUCINATION, + f"Schema hallucination: unknown table {table_name} not in catalog", risk="medium", - evidence={"table": table_name}, + evidence={"table": table_name, "known_tables": sorted(ctx.catalog.tables)}, ) if isinstance(stmt, exp.Insert): @@ -130,8 +130,11 @@ def _columns_and_types( if unknown: return GuardResult.block( NAME, - RULE_SCHEMA, - f"未知列 {unknown}，表 {spec.name} 的列为 {sorted(spec.columns)}", + RULE_HALLUCINATION, + ( + f"Schema hallucination: unknown column(s) {unknown} on table " + f"{spec.name}; known columns {sorted(spec.columns)}" + ), risk="medium", evidence={"unknown_columns": unknown, "table": spec.name}, ) @@ -160,3 +163,90 @@ def _columns_and_types( evidence={"not_allowed": not_allowed}, ) return None + + +def _cte_aliases(stmt) -> set[str]: + """Names introduced by WITH … AS (…); not catalog tables.""" + names: set[str] = set() + if stmt is None: + return names + from write_gate.idents import ident + + for cte in stmt.find_all(exp.CTE): + alias = getattr(cte, "alias_or_name", None) + if alias: + names.add(str(alias).lower()) + continue + alias_node = cte.args.get("alias") + if alias_node is not None: + name = ident(alias_node) or ident(getattr(alias_node, "this", None)) + if name: + names.add(name) + return names + + +def _check_select_hallucination(parsed, ctx) -> GuardResult: + """Block SELECT on unknown tables/columns (schema hallucination). + + CTE aliases are ignored. Honors ``allow_unknown_tables`` / + ``allow_unknown_columns`` policy knobs. ``SELECT *`` only validates tables. + """ + policy = ctx.policy + allow_unknown_tables = bool(getattr(policy, "allow_unknown_tables", False)) + allow_unknown_columns = bool(getattr(policy, "allow_unknown_columns", False)) + + catalog = ctx.catalog + tables = list(getattr(parsed, "tables_referenced", None) or []) + if parsed.table and parsed.table not in tables: + tables.insert(0, parsed.table) + + cte_names = _cte_aliases(parsed.statement) + physical = [t for t in tables if t not in cte_names] + + if not allow_unknown_tables: + unknown_tables = [t for t in physical if catalog.table(t) is None] + if unknown_tables: + return GuardResult.block( + NAME, + RULE_HALLUCINATION, + ( + f"Schema hallucination: unknown table(s) {unknown_tables} " + f"not in catalog allowlist {sorted(catalog.tables)}" + ), + risk="medium", + evidence={ + "unknown_tables": unknown_tables, + "known_tables": sorted(catalog.tables), + "cte_aliases": sorted(cte_names), + }, + ) + + if parsed.star or allow_unknown_columns: + return GuardResult.pass_(NAME) + + findings = getattr(parsed, "findings", None) + cols = list(getattr(findings, "columns", None) or parsed.select_columns or []) + if not cols or not physical: + return GuardResult.pass_(NAME) + + known: set[str] = set() + for tname in physical: + spec = catalog.table(tname) + if spec is not None: + known |= set(spec.columns.keys()) + if not known: + return GuardResult.pass_(NAME) + + unknown_cols = [c for c in cols if c not in known] + if unknown_cols: + return GuardResult.block( + NAME, + RULE_HALLUCINATION, + ( + f"Schema hallucination: unknown column(s) {unknown_cols} " + f"for tables {physical}; known columns {sorted(known)}" + ), + risk="medium", + evidence={"unknown_columns": unknown_cols, "tables": physical}, + ) + return GuardResult.pass_(NAME) diff --git a/src/write_gate/idents.py b/src/write_gate/idents.py new file mode 100644 index 0000000..c86a113 --- /dev/null +++ b/src/write_gate/idents.py @@ -0,0 +1,48 @@ +"""Shared SQL identifier / literal helpers (breaks parser ↔ ast_patterns cycle).""" + +from __future__ import annotations + +from typing import Any + +from sqlglot import exp + + +def ident(node: exp.Expression | None) -> str | None: + if node is None: + return None + if isinstance(node, exp.Table): + return node.name.lower() if node.name else None + if isinstance(node, exp.Schema): + return ident(node.this) + if isinstance(node, exp.Identifier): + return node.name.lower() + if isinstance(node, exp.Column): + return node.name.lower() if node.name else None + name = getattr(node, "name", None) + return str(name).lower() if name else None + + +def literal_value(node: exp.Expression | None) -> Any: + if node is None: + return None + if isinstance(node, exp.Null): + return None + if isinstance(node, exp.Cast): + return literal_value(node.this) + if isinstance(node, (exp.TsOrDsToDate, exp.Date)): + return literal_value(node.this) if node.this else node.sql() + if isinstance(node, exp.Literal): + raw = node.this + if node.is_int: + try: + return int(raw) + except (TypeError, ValueError): + return raw + if node.is_number: + try: + return float(raw) + except (TypeError, ValueError): + return raw + return str(raw) + sql = node.sql(dialect="duckdb").strip().strip("'\"") + return sql diff --git a/src/write_gate/mcp_server.py b/src/write_gate/mcp_server.py index ed3a26e..3160eff 100644 --- a/src/write_gate/mcp_server.py +++ b/src/write_gate/mcp_server.py @@ -21,7 +21,8 @@ def create_server( policy: str | None = None, agent: str = "mcp", ): - """Build a FastMCP server exposing query_sql and write_sql (ALLOW executes).""" + """Build a FastMCP server exposing query/write + DataPilot BLOCK/EXECUTE.""" + from write_gate.mcp_tools import datapilot_check, datapilot_execute from write_gate.mcp_tools import query_sql as check_query from write_gate.mcp_tools import write_sql as check_write @@ -45,6 +46,28 @@ def write_sql(sql: str) -> dict[str, Any]: """Evaluate INSERT/UPDATE/DELETE/DDL through sql-write-gate. ALLOW executes.""" return check_write(sql, **gate_kwargs) + @mcp.tool() + def datapilot_block_or_execute( + sql: str, + execute: bool = False, + actor: str | None = None, + model_id: str | None = None, + prompt_summary: str | None = None, + ) -> dict[str, Any]: + """Stable DataPilot API: returns datapilot=BLOCK|EXECUTE|APPROVAL. + + When execute=false (default), only evaluates. When execute=true, runs + SQL only if the gate returns ALLOW. + """ + fn = datapilot_execute if execute else datapilot_check + return fn( + sql, + actor=actor, + model_id=model_id, + prompt_summary=prompt_summary, + **gate_kwargs, + ) + return mcp diff --git a/src/write_gate/mcp_tools.py b/src/write_gate/mcp_tools.py index 1c62ff5..f92eab2 100644 --- a/src/write_gate/mcp_tools.py +++ b/src/write_gate/mcp_tools.py @@ -100,7 +100,10 @@ def decision_payload( "operation": decision.operation, "table": decision.table, "risk": decision.risk, + "risk_score": getattr(decision, "risk_score", 0), + "risk_factors": list(getattr(decision, "risk_factors", None) or []), "executed": bool(executed), + "product": "SQLGuard", } if getattr(decision, "approval_id", None): payload["approval_id"] = decision.approval_id @@ -207,3 +210,61 @@ def write_sql( approvals_path=approvals_path, agent=agent, ) + + +def datapilot_check( + sql: str, + *, + database: str | None = None, + db_path: str | Path | None = None, + catalog_path: str | Path | None = None, + policy_path: str | Path | None = None, + agent: str = "mcp", + actor: str | None = None, + model_id: str | None = None, + prompt_summary: str | None = None, +) -> dict[str, Any]: + """DataPilot BLOCK/EXECUTE check (never writes).""" + from write_gate.datapilot import block_or_execute + + return block_or_execute( + sql, + execute=False, + database=database, + db_path=str(db_path) if db_path else None, + catalog_path=str(catalog_path) if catalog_path else None, + policy_path=str(policy_path) if policy_path else None, + agent=agent, + actor=actor, + model_id=model_id, + prompt_summary=prompt_summary, + ) + + +def datapilot_execute( + sql: str, + *, + database: str | None = None, + db_path: str | Path | None = None, + catalog_path: str | Path | None = None, + policy_path: str | Path | None = None, + agent: str = "mcp", + actor: str | None = None, + model_id: str | None = None, + prompt_summary: str | None = None, +) -> dict[str, Any]: + """DataPilot EXECUTE path — runs SQL only when gate returns ALLOW.""" + from write_gate.datapilot import block_or_execute + + return block_or_execute( + sql, + execute=True, + database=database, + db_path=str(db_path) if db_path else None, + catalog_path=str(catalog_path) if catalog_path else None, + policy_path=str(policy_path) if policy_path else None, + agent=agent, + actor=actor, + model_id=model_id, + prompt_summary=prompt_summary, + ) diff --git a/src/write_gate/parser.py b/src/write_gate/parser.py index 9fd4897..8faf05b 100644 --- a/src/write_gate/parser.py +++ b/src/write_gate/parser.py @@ -10,7 +10,9 @@ import sqlglot from sqlglot import exp +from write_gate.ast_patterns import AstFindings, analyze_statement from write_gate.decision import RULE_SCHEMA, RULE_UNSUPPORTED +from write_gate.idents import ident, literal_value _DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$") @@ -57,48 +59,11 @@ class ParsedSQL: insert_rows: list[list[exp.Expression]] | None = None error: str | None = None error_rule: str = RULE_SCHEMA + findings: AstFindings = field(default_factory=AstFindings) + tables_referenced: list[str] = field(default_factory=list) + dangerous_flags: list[str] = field(default_factory=list) -def ident(node: exp.Expression | None) -> str | None: - if node is None: - return None - if isinstance(node, exp.Table): - return node.name.lower() if node.name else None - if isinstance(node, exp.Schema): - return ident(node.this) - if isinstance(node, exp.Identifier): - return node.name.lower() - if isinstance(node, exp.Column): - return node.name.lower() if node.name else None - name = getattr(node, "name", None) - return str(name).lower() if name else None - - -def literal_value(node: exp.Expression | None) -> Any: - if node is None: - return None - if isinstance(node, exp.Null): - return None - if isinstance(node, exp.Cast): - return literal_value(node.this) - if isinstance(node, (exp.TsOrDsToDate, exp.Date)): - return literal_value(node.this) if node.this else node.sql() - if isinstance(node, exp.Literal): - raw = node.this - if node.is_int: - try: - return int(raw) - except (TypeError, ValueError): - return raw - if node.is_number: - try: - return float(raw) - except (TypeError, ValueError): - return raw - return str(raw) - sql = node.sql(dialect="duckdb").strip().strip("'\"") - return sql - def expected_kind(col_type: str) -> str: t = col_type.upper() @@ -915,4 +880,18 @@ def parse(sql: str, dialect: str = "duckdb") -> ParsedSQL: parsed.select_columns = cols parsed.columns = cols parsed.star = star + + # Stronger AST analysis (joins, tautology WHERE, referenced objects). + findings = analyze_statement( + stmt, + operation=parsed.operation, + has_where=parsed.has_where, + where=parsed.where, + ) + parsed.findings = findings + parsed.tables_referenced = list(findings.tables) + parsed.dangerous_flags = list(findings.dangerous_flags) + # Treat tautology WHERE as missing for downstream destructive guards. + if findings.tautology_where and parsed.operation in {"update", "delete"}: + parsed.has_where = False return parsed diff --git a/src/write_gate/registry.py b/src/write_gate/registry.py new file mode 100644 index 0000000..6f4f529 --- /dev/null +++ b/src/write_gate/registry.py @@ -0,0 +1,98 @@ +"""Pluggable guard / rule registry — extend the engine without forking it.""" + +from __future__ import annotations + +from typing import Callable, Iterable + +from write_gate.decision import GuardResult + +GuardFn = Callable[["Context"], GuardResult] # type: ignore[name-defined] + + +class GuardRegistry: + """Ordered registry of guard callables. + + Built-in guards register at import time via ``default_registry()``. + Third-party / site rules call ``register`` / ``unregister`` without + editing ``engine.py``. + """ + + def __init__(self) -> None: + self._guards: list[tuple[str, GuardFn]] = [] + + def register(self, name: str, fn: GuardFn, *, before: str | None = None) -> None: + """Register or replace a guard by name. + + If ``before`` is set, insert ahead of that named guard; otherwise append. + """ + self.unregister(name) + entry = (name, fn) + if before: + for i, (n, _) in enumerate(self._guards): + if n == before: + self._guards.insert(i, entry) + return + self._guards.append(entry) + + def unregister(self, name: str) -> bool: + before = len(self._guards) + self._guards = [(n, f) for n, f in self._guards if n != name] + return len(self._guards) < before + + def names(self) -> list[str]: + return [n for n, _ in self._guards] + + def functions(self) -> list[GuardFn]: + return [f for _, f in self._guards] + + def run(self, ctx) -> list[GuardResult]: + return [fn(ctx) for _, fn in self._guards] + + def clear(self) -> None: + self._guards.clear() + + def extend(self, items: Iterable[tuple[str, GuardFn]]) -> None: + for name, fn in items: + self.register(name, fn) + + +_DEFAULT: GuardRegistry | None = None + + +def default_registry() -> GuardRegistry: + """Singleton registry preloaded with built-in SQLGuard guards.""" + global _DEFAULT + if _DEFAULT is not None: + return _DEFAULT + from write_gate.guards import ( + check_ast_patterns, + check_blast_radius, + check_destructive, + check_environment, + check_explain_cost, + check_freshness, + check_permissions, + check_pii, + check_schema, + ) + + reg = GuardRegistry() + # Order matters: specific danger first, environment last. + reg.register("destructive", check_destructive) + reg.register("ast_patterns", check_ast_patterns) + reg.register("schema", check_schema) + reg.register("permissions", check_permissions) + reg.register("pii", check_pii) + reg.register("freshness", check_freshness) + reg.register("blast_radius", check_blast_radius) + reg.register("explain_cost", check_explain_cost) + reg.register("environment", check_environment) + _DEFAULT = reg + return reg + + +def reset_default_registry() -> GuardRegistry: + """Drop the singleton (tests). Next ``default_registry()`` rebuilds.""" + global _DEFAULT + _DEFAULT = None + return default_registry() diff --git a/src/write_gate/risk.py b/src/write_gate/risk.py new file mode 100644 index 0000000..9c9407b --- /dev/null +++ b/src/write_gate/risk.py @@ -0,0 +1,166 @@ +"""Numeric risk score (0–100) with contributing factors for SQLGuard Decisions.""" + +from __future__ import annotations + +from typing import Any, Iterable + +from write_gate.decision import ( + ACTION_ALLOW, + ACTION_APPROVAL, + ACTION_BLOCK, + RISK_CRITICAL, + RISK_MEDIUM, + VERDICT_APPROVAL, + VERDICT_BLOCK, + GuardResult, +) + +# Base weights (capped later). Factors are additive then clamped to 0..100. +_FACTOR_WEIGHTS: dict[str, int] = { + "block": 80, + "approval": 35, + "ddl": 50, + "delete": 40, + "update": 25, + "insert": 10, + "missing_where": 45, + "tautology_where": 45, + "full_table_write": 50, + "cartesian_join": 40, + "pii": 30, + "restricted": 50, + "expired_partition": 35, + "blast_radius": 40, + "blast_unknown": 45, + "schema_hallucination": 55, + "permission_denied": 60, + "explain_cost_high": 25, + "environment_block": 50, + "unsupported_sql": 55, +} + + +def _add(factors: list[str], name: str) -> None: + if name not in factors: + factors.append(name) + + +def score_from_results( + *, + action: str, + operation: str | None, + results: Iterable[GuardResult], + dangerous_flags: Iterable[str] | None = None, + explain_cost: float | None = None, + explain_cost_threshold: float | None = None, +) -> tuple[int, list[str]]: + """Compute (risk_score, risk_factors) from guard results + AST flags.""" + factors: list[str] = [] + total = 0 + + if action == ACTION_BLOCK: + _add(factors, "block") + total += _FACTOR_WEIGHTS["block"] + elif action == ACTION_APPROVAL: + _add(factors, "approval") + total += _FACTOR_WEIGHTS["approval"] + + op = (operation or "").lower() + if op in _FACTOR_WEIGHTS: + _add(factors, op) + total += _FACTOR_WEIGHTS[op] + + for flag in dangerous_flags or []: + key = str(flag) + if key in _FACTOR_WEIGHTS: + _add(factors, key) + total += _FACTOR_WEIGHTS[key] + + for result in results: + rid = result.rule_id or "" + if result.verdict == VERDICT_BLOCK: + if rid in {"delete_without_where", "update_without_where", "full_table_write"}: + _add(factors, "missing_where") + total += _FACTOR_WEIGHTS["missing_where"] + elif rid == "cartesian_join": + _add(factors, "cartesian_join") + total += _FACTOR_WEIGHTS["cartesian_join"] + elif rid in {"schema_hallucination", "schema_mismatch"}: + if "unknown" in (result.reason or "").lower() or rid == "schema_hallucination": + _add(factors, "schema_hallucination") + total += _FACTOR_WEIGHTS["schema_hallucination"] + elif rid == "table_permission": + _add(factors, "permission_denied") + total += _FACTOR_WEIGHTS["permission_denied"] + elif rid == "pii_column": + _add(factors, "pii") + total += _FACTOR_WEIGHTS["pii"] + elif rid == "restricted_column": + _add(factors, "restricted") + total += _FACTOR_WEIGHTS["restricted"] + elif rid == "expired_partition": + _add(factors, "expired_partition") + total += _FACTOR_WEIGHTS["expired_partition"] + elif rid == "blast_radius_exceeded": + _add(factors, "blast_radius") + total += _FACTOR_WEIGHTS["blast_radius"] + elif rid == "blast_radius_unknown": + _add(factors, "blast_unknown") + total += _FACTOR_WEIGHTS["blast_unknown"] + elif rid == "environment_policy": + _add(factors, "environment_block") + total += _FACTOR_WEIGHTS["environment_block"] + elif rid == "unsupported_sql": + _add(factors, "unsupported_sql") + total += _FACTOR_WEIGHTS["unsupported_sql"] + elif rid in {"drop_table", "truncate_table", "alter_table"}: + _add(factors, "ddl") + total += _FACTOR_WEIGHTS["ddl"] + elif result.verdict == VERDICT_APPROVAL: + if "approval" not in factors: + _add(factors, "approval") + total += 15 # smaller bump when already counted + + if ( + explain_cost is not None + and explain_cost_threshold is not None + and explain_cost >= explain_cost_threshold + ): + _add(factors, "explain_cost_high") + total += _FACTOR_WEIGHTS["explain_cost_high"] + + score = max(0, min(100, int(total))) + # Ensure BLOCK is never scored trivially low + if action == ACTION_BLOCK and score < 50: + score = 50 + if action == ACTION_ALLOW and not factors: + score = 5 if op in {"insert", "update", "delete"} else 0 + return score, factors + + +def qualitative_from_score(score: int) -> str: + if score >= 70: + return RISK_CRITICAL + if score >= 35: + return RISK_MEDIUM + return "low" + + +def attach_explain_evidence( + evidence: dict[str, Any], + *, + cost: float | None, + plan: str | None, + skipped: str | None, +) -> dict[str, Any]: + blob: dict[str, Any] = dict(evidence) + explain: dict[str, Any] = {} + if cost is not None: + explain["cost"] = cost + if plan is not None: + explain["plan"] = plan + if skipped is not None: + explain["skipped"] = skipped + if explain: + blob["explain"] = explain + return blob diff --git a/src/write_gate/sqlguard.py b/src/write_gate/sqlguard.py new file mode 100644 index 0000000..fe62a2a --- /dev/null +++ b/src/write_gate/sqlguard.py @@ -0,0 +1,28 @@ +"""Optional SQLGuard product alias over write_gate. + +DataPilot and docs may ``import sqlguard`` / ``from write_gate import sqlguard``. +Package name and CLI remain ``write_gate`` / ``sql-write-gate``. +""" + +from __future__ import annotations + +from write_gate import Decision, Evidence, WriteGate, __version__ +from write_gate.api import decision_response, handle_datapilot_request, serve +from write_gate.engine import evaluate +from write_gate.policy import evaluate as policy_evaluate + +PRODUCT = "SQLGuard" +VERSION = __version__ + +__all__ = [ + "PRODUCT", + "VERSION", + "WriteGate", + "Decision", + "Evidence", + "evaluate", + "policy_evaluate", + "decision_response", + "handle_datapilot_request", + "serve", +] diff --git a/src/write_gate/templates/policy.yaml b/src/write_gate/templates/policy.yaml index 197fe22..d92a3f4 100644 --- a/src/write_gate/templates/policy.yaml +++ b/src/write_gate/templates/policy.yaml @@ -10,3 +10,11 @@ rules: limits: update_rows: 100 delete_rows: 50 +# SQLGuard — optional GameStream-style table permissions +# permissions: +# enforce: true +# tables: +# orders: [select, insert, update] +# hallucination: +# allow_unknown_tables: false +# allow_unknown_columns: false diff --git a/src/write_gate/wrapper.py b/src/write_gate/wrapper.py index 4760277..e01e255 100644 --- a/src/write_gate/wrapper.py +++ b/src/write_gate/wrapper.py @@ -66,6 +66,9 @@ def __init__( agent: str = "cli", database: str | None = None, database_url: str | None = None, + actor: str | None = None, + model_id: str | None = None, + prompt_summary: str | None = None, ) -> None: backend, target = resolve_target( database=database, @@ -84,6 +87,9 @@ def __init__( Path(approvals_path) if approvals_path else default_approvals_path() ) self.agent = agent + self.actor = actor or agent + self.model_id = model_id + self.prompt_summary = prompt_summary self.database_config_id = None self._conn = conn self._owns_conn = conn is None @@ -253,6 +259,8 @@ def _audit( executed: bool | None = None, execution_outcome: str | None = None, error_class: str | None = None, + latency_ms: float | None = None, + success: bool | None = None, ) -> None: append_audit( decision, @@ -264,10 +272,18 @@ def _audit( execution_outcome=execution_outcome, error_class=error_class, request_id=self.request_id, + actor=self.actor, + model_id=self.model_id, + prompt_summary=self.prompt_summary, + latency_ms=latency_ms, + success=success, ) def check(self, sql: str) -> Decision: + import time as _time + # Apply statement timeout to evaluate (blast-radius COUNT, etc.). + t0 = _time.perf_counter() if self.runtime.timeout_enabled: try: decision = run_with_timeout( @@ -289,11 +305,17 @@ def check(self, sql: str) -> Decision: executed=False, execution_outcome="failed", error_class=type(exc).__name__, + latency_ms=(_time.perf_counter() - t0) * 1000.0, + success=False, ) return decision else: decision = self._evaluate(sql, use_conn=False) - self._audit(decision) + self._audit( + decision, + latency_ms=(_time.perf_counter() - t0) * 1000.0, + success=decision.action == ACTION_ALLOW, + ) return decision def execute(self, sql: str) -> tuple[Decision, Any]: @@ -302,6 +324,9 @@ def execute(self, sql: str) -> tuple[Decision, Any]: REQUIRE_APPROVAL is enqueued and not executed. BLOCK is not queued and not executed. check() stays evaluate-only (no enqueue). """ + import time as _time + + t0 = _time.perf_counter() decision = self._evaluate(sql, use_conn=True) if decision.action == ACTION_APPROVAL: rec = enqueue_approval( @@ -321,6 +346,8 @@ def execute(self, sql: str) -> tuple[Decision, Any]: decision, executed=False, execution_outcome="queued", + latency_ms=(_time.perf_counter() - t0) * 1000.0, + success=False, ) return decision, None if decision.action != ACTION_ALLOW: @@ -328,6 +355,8 @@ def execute(self, sql: str) -> tuple[Decision, Any]: decision, executed=False, execution_outcome="blocked", + latency_ms=(_time.perf_counter() - t0) * 1000.0, + success=False, ) return decision, None try: @@ -341,12 +370,16 @@ def execute(self, sql: str) -> tuple[Decision, Any]: executed=False, execution_outcome=exec_outcome, error_class=type(exc).__name__, + latency_ms=(_time.perf_counter() - t0) * 1000.0, + success=False, ) raise self._audit( decision, executed=True, execution_outcome="executed", + latency_ms=(_time.perf_counter() - t0) * 1000.0, + success=True, ) return decision, result diff --git a/tests/test_demo_cases.py b/tests/test_demo_cases.py index 8cf631a..d22949a 100644 --- a/tests/test_demo_cases.py +++ b/tests/test_demo_cases.py @@ -71,7 +71,7 @@ def test_schema_mismatch_rejected_no_write(gate): ev, result = gate.execute(SCHEMA_MISMATCH_SQL) after = gate.conn.execute("SELECT COUNT(*) FROM orders").fetchone()[0] assert ev.allowed is False - assert ev.rule_id == "schema_mismatch" + assert ev.rule_id == "schema_hallucination" assert result is None assert after == before diff --git a/tests/test_policy.py b/tests/test_policy.py index 86abbd5..0617a15 100644 --- a/tests/test_policy.py +++ b/tests/test_policy.py @@ -47,7 +47,7 @@ def test_pii_insert_rejected(): def test_schema_unknown_column_rejected(): ev = evaluate(SCHEMA_MISMATCH_SQL, _catalog()) assert ev.allowed is False - assert ev.rule_id == "schema_mismatch" + assert ev.rule_id == "schema_hallucination" assert "not_a_column" in ev.message @@ -63,7 +63,7 @@ def test_unknown_table_rejected(): _catalog(), ) assert ev.allowed is False - assert ev.rule_id == "schema_mismatch" + assert ev.rule_id == "schema_hallucination" def test_pii_update_rejected(): @@ -92,7 +92,7 @@ def test_select_star_requires_approval_for_pii(): def test_select_pii_column_requires_approval(): - ev = evaluate("SELECT id, email FROM orders LIMIT 10", _catalog()) + ev = evaluate("SELECT order_id, email FROM orders LIMIT 10", _catalog()) assert ev.action == "REQUIRE_APPROVAL" assert ev.rule_id == "pii_column" assert ev.allowed is False diff --git a/tests/test_v101.py b/tests/test_v101.py index eaa2bfc..2b26dd6 100644 --- a/tests/test_v101.py +++ b/tests/test_v101.py @@ -36,11 +36,11 @@ def fetchmany(self, n): def test_version_is_101(): - assert __version__ == "1.0.1" - text = (ROOT / "pyproject.toml").read_text(encoding="utf-8") - assert 'version = "1.0.1"' in text - init = (ROOT / "src" / "write_gate" / "__init__.py").read_text(encoding="utf-8") - assert '__version__ = "1.0.1"' in init + # Pinned suite for the 1.0.1 fixes; current tree may be newer (1.1.x). + parts = [int(x) for x in __version__.split(".")[:3]] + assert parts >= [1, 0, 1], __version__ + text = (ROOT / "CHANGELOG.md").read_text(encoding="utf-8") + assert "## [1.0.1]" in text def test_changelog_has_101(): diff --git a/tests/test_v110.py b/tests/test_v110.py new file mode 100644 index 0000000..bdb0e4d --- /dev/null +++ b/tests/test_v110.py @@ -0,0 +1,244 @@ +"""SQLGuard 1.1.0: AST patterns, permissions, risk score, hallucination, audit, API.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from write_gate import WriteGate, __version__ +from write_gate.api import handle_datapilot_request +from write_gate.audit import read_audit +from write_gate.catalog import load_catalog +from write_gate.config import demo_policy, policy_from_dict +from write_gate.parser import parse +from write_gate.policy import evaluate +from write_gate.sqlguard import PRODUCT + +ROOT = Path(__file__).resolve().parents[1] + + +def test_version_is_110(): + assert __version__ == "1.1.0" + assert 'version = "1.1.0"' in (ROOT / "pyproject.toml").read_text(encoding="utf-8") + assert "## [1.1.0]" in (ROOT / "CHANGELOG.md").read_text(encoding="utf-8") + assert PRODUCT == "SQLGuard" + + +def test_ast_detects_cartesian_join(): + parsed = parse("SELECT * FROM orders CROSS JOIN orders") + assert parsed.findings.cartesian_joins + assert "cartesian_join" in parsed.dangerous_flags + + +def test_ast_detects_comma_join_cartesian(): + parsed = parse("SELECT * FROM orders, orders") + assert any(j.is_cartesian for j in parsed.findings.joins) + + +def test_cartesian_join_blocked(): + # Use two aliases of known table to avoid hallucination on second name + sql = "SELECT o.order_id FROM orders o CROSS JOIN orders p" + ev = evaluate(sql, load_catalog()) + assert ev.action == "BLOCK" + assert ev.rule_id == "cartesian_join" + assert ev.risk_score >= 50 + + +def test_tautology_where_treated_as_full_table_delete(): + ev = evaluate("DELETE FROM orders WHERE 1=1", load_catalog()) + assert ev.action == "BLOCK" + assert ev.rule_id == "delete_without_where" + + +def test_tautology_where_true_update(): + ev = evaluate("UPDATE orders SET status = 'x' WHERE TRUE", load_catalog()) + assert ev.action == "BLOCK" + assert ev.rule_id == "update_without_where" + + +def test_schema_hallucination_unknown_table_select(): + ev = evaluate("SELECT * FROM game_sessions_not_real", load_catalog()) + assert ev.action == "BLOCK" + assert ev.rule_id == "schema_hallucination" + schema_ev = ev.evidence.get("schema") or ev.evidence + assert "unknown_tables" in schema_ev or "game_sessions" in ev.reason + + +def test_schema_hallucination_unknown_column_select(): + ev = evaluate("SELECT not_a_real_col FROM orders", load_catalog()) + assert ev.action == "BLOCK" + assert ev.rule_id == "schema_hallucination" + + +def test_schema_hallucination_unknown_column_insert(): + sql = ( + "INSERT INTO orders (order_id, user_id, amount, dt, status, not_a_column) " + "VALUES (1, 1, 1.0, '2026-09-01', 'paid', 1)" + ) + ev = evaluate(sql, load_catalog()) + assert ev.action == "BLOCK" + assert ev.rule_id == "schema_hallucination" + + +def test_permissions_deny_table_op(): + policy = policy_from_dict( + { + "environment": "demo", + "rules": {"select": "allow", "insert": "allow", "update": "allow", "delete": "allow", "ddl": "block"}, + "permissions": {"enforce": True, "tables": {"orders": ["select"]}}, + } + ) + sql = ( + "INSERT INTO orders (order_id, user_id, amount, dt, status) " + "VALUES (910001, 1, 1.0, '2026-09-01', 'paid')" + ) + ev = evaluate(sql, load_catalog(), policy=policy) + assert ev.action == "BLOCK" + assert ev.rule_id == "table_permission" + + +def test_permissions_allow_listed_op(): + policy = policy_from_dict( + { + "environment": "demo", + "rules": {"select": "allow", "insert": "allow", "update": "allow", "delete": "allow", "ddl": "block"}, + "permissions": {"enforce": True, "tables": {"orders": ["select", "insert"]}}, + } + ) + sql = ( + "INSERT INTO orders (order_id, user_id, amount, dt, status) " + "VALUES (910002, 1, 1.0, '2026-09-01', 'paid')" + ) + ev = evaluate(sql, load_catalog(), policy=policy) + assert ev.action == "ALLOW" + assert isinstance(ev.risk_score, int) + assert 0 <= ev.risk_score <= 100 + + +def test_risk_score_on_decision_dict(): + ev = evaluate("DELETE FROM orders", load_catalog()) + d = ev.to_dict() + assert "risk_score" in d + assert "risk_factors" in d + assert d["risk_score"] >= 50 + assert "block" in d["risk_factors"] or "missing_where" in d["risk_factors"] or "delete" in d["risk_factors"] + + +def test_richer_audit_fields(tmp_path): + audit = tmp_path / "audit.jsonl" + gate = WriteGate( + db_path=tmp_path / "wh.duckdb", + policy=demo_policy(), + audit_path=audit, + agent="test", + actor="datapilot-bot", + model_id="gpt-test", + prompt_summary="delete all orders", + ) + gate.check("DELETE FROM orders") + rows = read_audit(audit, limit=5) + assert len(rows) == 1 + rec = rows[0] + assert rec["actor"] == "datapilot-bot" + assert rec["model_id"] == "gpt-test" + assert rec["prompt_summary"] == "delete all orders" + assert "risk_score" in rec + assert "latency_ms" in rec + assert rec["success"] is False + assert rec["decision"] == "BLOCK" + + +def test_datapilot_api_check_and_execute(tmp_path): + # Use isolated db copy of seed if present + defaults = { + "db_path": str(ROOT / "seed" / "warehouse.duckdb"), + "policy": str(ROOT / "examples" / "policy.demo.yaml"), + "agent": "datapilot", + } + status, payload = handle_datapilot_request( + "POST", + "/v1/check", + {"sql": "DELETE FROM orders", "actor": "pilot", "model_id": "m1"}, + defaults=defaults, + ) + assert status == 200 + assert payload["action"] == "BLOCK" + assert payload["executed"] is False + assert payload["product"] == "SQLGuard" + assert payload["risk_score"] >= 50 + + status, health = handle_datapilot_request("GET", "/healthz", None, defaults=defaults) + assert status == 200 + assert health["ok"] is True + assert health["version"] == "1.1.0" + + # Legal insert — use check (no mutate shared seed); execute path covered by demo + status, payload = handle_datapilot_request( + "POST", + "/v1/check", + { + "sql": ( + "INSERT INTO orders (order_id, user_id, amount, dt, status) " + "VALUES (910099, 1, 1.0, '2026-09-01', 'paid')" + ), + "actor": "pilot", + }, + defaults=defaults, + ) + assert status == 200 + assert payload["action"] == "ALLOW" + assert payload["executed"] is False + assert payload["risk_score"] >= 0 + + +def test_explain_degrades_offline(): + from write_gate.explain import estimate_cost + + cost, plan, skip = estimate_cost(None, "SELECT 1") + assert cost is None + assert skip == "no connection" + + +def test_registry_can_register_custom_guard(): + from write_gate.decision import GuardResult, VERDICT_BLOCK + from write_gate.registry import GuardRegistry, default_registry + + reg = GuardRegistry() + for name, fn in zip(default_registry().names(), default_registry().functions()): + reg.register(name, fn) + + def boom(ctx): + return GuardResult.block("custom", "custom_rule", "custom blocked") + + reg.register("custom", boom, before="environment") + assert "custom" in reg.names() + from write_gate.engine import evaluate + from write_gate.catalog import load_catalog + from write_gate.config import demo_policy + + ev = evaluate("SELECT order_id FROM orders", load_catalog(), policy=demo_policy(), registry=reg) + assert ev.action == "BLOCK" + assert ev.rule_id == "custom_rule" + + +def test_ddl_drop_sets_dangerous_flags(): + parsed = parse("DROP TABLE orders") + assert "ddl" in parsed.dangerous_flags or parsed.operation == "ddl" + ev = evaluate("DROP TABLE orders", load_catalog()) + assert ev.action == "BLOCK" + assert ev.rule_id == "drop_table" + + +def test_datapilot_field_on_api_response(): + defaults = { + "db_path": str(ROOT / "seed" / "warehouse.duckdb"), + "policy": str(ROOT / "examples" / "policy.demo.yaml"), + } + status, payload = handle_datapilot_request( + "POST", "/v1/check", {"sql": "DELETE FROM orders"}, defaults=defaults + ) + assert status == 200 + assert payload.get("datapilot") == "BLOCK" +