sql-write-gate 0.16.0__py3-none-any.whl

write_gate/parser.py

@@ -0,0 +1,367 @@
+"""Parse a single SQL statement into a structured form (sqlglot AST, no LLM)."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from datetime import date, datetime
+from typing import Any
+
+import sqlglot
+from sqlglot import exp
+
+from write_gate.decision import RULE_SCHEMA
+
+_DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$")
+
+WRITE_TYPE_NAMES = (
+    "Insert",
+    "Update",
+    "Delete",
+    "Merge",
+    "Create",
+    "Drop",
+    "Alter",
+    "AlterTable",
+    "Command",
+    "Copy",
+    "TruncateTable",
+    "Replace",
+)
+
+DDL_TYPE_NAMES = (
+    "Create",
+    "Drop",
+    "Alter",
+    "AlterTable",
+    "TruncateTable",
+    "Command",
+)
+
+
+@dataclass
+class ParsedSQL:
+    sql: str
+    statement: exp.Expression | None = None
+    operation: str = "unknown"  # select | insert | update | delete | ddl | unknown
+    table: str | None = None
+    columns: list[str] = field(default_factory=list)
+    write_columns: list[str] = field(default_factory=list)
+    select_columns: list[str] = field(default_factory=list)
+    star: bool = False
+    where: exp.Expression | None = None
+    has_where: bool = False
+    assignments: dict[str, exp.Expression | None] = field(default_factory=dict)
+    insert_rows: list[list[exp.Expression]] | None = None
+    error: str | None = None
+    error_rule: str = RULE_SCHEMA
+
+
+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()
+    if t in {"INTEGER", "INT", "BIGINT", "SMALLINT", "TINYINT"}:
+        return "INTEGER"
+    if t in {"DOUBLE", "FLOAT", "REAL", "DECIMAL", "NUMERIC"}:
+        return "DOUBLE"
+    if t in {"DATE", "TIMESTAMP"}:
+        return "DATE"
+    return "VARCHAR"
+
+
+def type_ok(expected: str, node: exp.Expression | None) -> bool:
+    if node is None or isinstance(node, exp.Null):
+        return True
+    kind = expected_kind(expected)
+    value = literal_value(node)
+    if kind == "INTEGER":
+        return isinstance(value, int) and not isinstance(value, bool)
+    if kind == "DOUBLE":
+        return isinstance(value, (int, float)) and not isinstance(value, bool)
+    if kind == "DATE":
+        text = str(value)
+        if not _DATE_RE.match(text):
+            return False
+        try:
+            date.fromisoformat(text)
+            return True
+        except ValueError:
+            return False
+    return True
+
+
+def parse_date(value: Any) -> date | None:
+    if value is None:
+        return None
+    if isinstance(value, date) and not isinstance(value, datetime):
+        return value
+    text = str(value)
+    if not _DATE_RE.match(text):
+        return None
+    try:
+        return date.fromisoformat(text)
+    except ValueError:
+        return None
+
+
+def _write_types() -> tuple[type, ...]:
+    return tuple(getattr(exp, n) for n in WRITE_TYPE_NAMES if hasattr(exp, n))
+
+
+def _ddl_types() -> tuple[type, ...]:
+    return tuple(getattr(exp, n) for n in DDL_TYPE_NAMES if hasattr(exp, n))
+
+
+def is_read_only(stmt: exp.Expression) -> bool:
+    write_types = _write_types()
+    if write_types and isinstance(stmt, write_types):
+        return False
+    if isinstance(stmt, (exp.Select, exp.Union, exp.Except, exp.Intersect)):
+        return True
+    return isinstance(stmt, exp.Query) and not isinstance(stmt, write_types)
+
+
+def classify_operation(stmt: exp.Expression) -> str:
+    if isinstance(stmt, exp.Insert):
+        return "insert"
+    if isinstance(stmt, exp.Update):
+        return "update"
+    if isinstance(stmt, exp.Delete):
+        return "delete"
+    ddl_types = _ddl_types()
+    if ddl_types and isinstance(stmt, ddl_types):
+        return "ddl"
+    if isinstance(stmt, exp.Merge):
+        return "ddl"
+    if is_read_only(stmt):
+        return "select"
+    return "ddl"
+
+
+def _table_from_select(stmt: exp.Expression) -> str | None:
+    from_ = stmt.args.get("from_") or stmt.args.get("from")
+    if from_ is None:
+        return None
+    this = from_.this if isinstance(from_, exp.From) else from_
+    return ident(this)
+
+
+def _table_from_truncate(stmt: exp.Expression) -> str | None:
+    for item in stmt.expressions or []:
+        name = ident(item)
+        if name:
+            return name
+    return ident(stmt.this)
+
+
+def extract_table(stmt: exp.Expression) -> str | None:
+    if isinstance(stmt, exp.TruncateTable) or type(stmt).__name__ == "TruncateTable":
+        return _table_from_truncate(stmt)
+    if isinstance(stmt, (exp.Select, exp.Union, exp.Except, exp.Intersect)):
+        return _table_from_select(stmt)
+    if isinstance(stmt, exp.Query):
+        return _table_from_select(stmt)
+    return ident(stmt.this)
+
+
+def extract_where(stmt: exp.Expression) -> exp.Expression | None:
+    where = stmt.args.get("where")
+    if where is None:
+        return None
+    if isinstance(where, exp.Where):
+        return where.this
+    return where
+
+
+def where_sql(where: exp.Expression | None, dialect: str = "duckdb") -> str | None:
+    if where is None:
+        return None
+    if isinstance(where, exp.Where):
+        return where.this.sql(dialect=dialect) if where.this else None
+    return where.sql(dialect=dialect)
+
+
+def _select_columns(stmt: exp.Expression) -> tuple[list[str], bool]:
+    cols: list[str] = []
+    star = False
+    expressions = stmt.expressions or []
+    for item in expressions:
+        if isinstance(item, exp.Star):
+            star = True
+            continue
+        if isinstance(item, exp.Alias):
+            name = ident(item.this) or ident(item)
+        else:
+            name = ident(item)
+        if name and name != "*":
+            cols.append(name)
+        elif isinstance(item, exp.Star) or (name == "*"):
+            star = True
+    return cols, star
+
+
+def insert_columns(stmt: exp.Insert, fallback: list[str] | None = None) -> list[str]:
+    target = stmt.this
+    if isinstance(target, exp.Schema) and target.expressions:
+        return [ident(c) or "" for c in target.expressions]
+    return list(fallback or [])
+
+
+def insert_rows(stmt: exp.Insert) -> list[list[exp.Expression]] | None:
+    values = stmt.expression
+    if isinstance(values, exp.Values):
+        rows: list[list[exp.Expression]] = []
+        for tup in values.expressions:
+            if isinstance(tup, exp.Tuple):
+                rows.append(list(tup.expressions))
+            else:
+                rows.append([tup])
+        return rows
+    return None
+
+
+def update_assignments(stmt: exp.Update) -> dict[str, exp.Expression | None]:
+    out: dict[str, exp.Expression | None] = {}
+    for item in stmt.expressions:
+        if isinstance(item, exp.EQ):
+            name = ident(item.this)
+            if name:
+                out[name] = item.expression
+    return out
+
+
+def partition_dates_from_where(
+    where: exp.Expression | None, part: str | None
+) -> list[date | None]:
+    if where is None or not part:
+        return []
+    dates: list[date | None] = []
+
+    def visit(node: exp.Expression) -> None:
+        if isinstance(node, exp.In) and ident(node.this) == part:
+            for item in node.expressions:
+                dates.append(parse_date(literal_value(item)))
+            return
+        if isinstance(node, exp.EQ):
+            left, right = node.this, node.expression
+            if ident(left) == part:
+                dates.append(parse_date(literal_value(right)))
+            elif ident(right) == part:
+                dates.append(parse_date(literal_value(left)))
+            return
+        for child in node.iter_expressions():
+            visit(child)
+
+    visit(where)
+    return dates
+
+
+def partition_dates_from_insert(
+    cols: list[str], rows: list[list[exp.Expression]], part: str | None
+) -> list[date | None]:
+    if not part:
+        return []
+    dates: list[date | None] = []
+    for row in rows:
+        assignments = dict(zip(cols, row))
+        if part not in assignments:
+            dates.append(None)
+        else:
+            dates.append(parse_date(literal_value(assignments[part])))
+    return dates
+
+
+def parse(sql: str, dialect: str = "duckdb") -> ParsedSQL:
+    original = sql
+    stripped = sql.strip().rstrip(";").strip()
+    parsed = ParsedSQL(sql=original)
+    if not stripped:
+        parsed.error = "SQL 为空，无法执行"
+        return parsed
+
+    read_dialect = dialect or "duckdb"
+    try:
+        statements = sqlglot.parse(stripped, read=read_dialect)
+    except sqlglot.errors.ParseError as exc:
+        if read_dialect != "duckdb":
+            try:
+                statements = sqlglot.parse(stripped, read="duckdb")
+            except sqlglot.errors.ParseError as exc2:
+                parsed.error = f"SQL 无法解析: {exc2}"
+                return parsed
+        else:
+            parsed.error = f"SQL 无法解析: {exc}"
+            return parsed
+
+    statements = [s for s in statements if s is not None]
+    if not statements:
+        parsed.error = "SQL 无法解析为空语句"
+        return parsed
+    if len(statements) != 1:
+        parsed.error = f"一次只允许一条语句，收到 {len(statements)} 条"
+        return parsed
+
+    stmt = statements[0]
+    parsed.statement = stmt
+    parsed.operation = classify_operation(stmt)
+    parsed.table = extract_table(stmt)
+    parsed.where = extract_where(stmt)
+    parsed.has_where = parsed.where is not None
+
+    if isinstance(stmt, exp.Insert):
+        cols = insert_columns(stmt)
+        parsed.write_columns = cols
+        parsed.columns = list(cols)
+        parsed.insert_rows = insert_rows(stmt)
+        parsed.assignments = {}
+    elif isinstance(stmt, exp.Update):
+        assignments = update_assignments(stmt)
+        parsed.assignments = assignments
+        parsed.write_columns = list(assignments.keys())
+        parsed.columns = list(assignments.keys())
+    elif is_read_only(stmt):
+        cols, star = _select_columns(stmt)
+        parsed.select_columns = cols
+        parsed.columns = cols
+        parsed.star = star
+    return parsed

write_gate/paths.py

@@ -0,0 +1,20 @@
+"""Project paths. Warehouse and catalog live under seed/."""
+
+from pathlib import Path
+
+PACKAGE_DIR = Path(__file__).resolve().parent
+PROJECT_ROOT = PACKAGE_DIR.parents[1]
+SEED_DIR = PROJECT_ROOT / "seed"
+EXAMPLES_DIR = PROJECT_ROOT / "examples"
+LOG_DIR = PROJECT_ROOT / ".logs"
+
+DB_PATH = SEED_DIR / "warehouse.duckdb"
+CATALOG_PATH = SEED_DIR / "catalog.json"
+ORDERS_CSV = SEED_DIR / "orders.csv"
+
+POLICY_PATH = PROJECT_ROOT / "policy.yaml"
+EXAMPLES_POLICY_PATH = EXAMPLES_DIR / "policy.yaml"
+DEMO_POLICY_PATH = EXAMPLES_DIR / "policy.demo.yaml"
+EXAMPLES_CATALOG_PATH = EXAMPLES_DIR / "catalog.json"
+AUDIT_PATH = LOG_DIR / "audit.jsonl"
+APPROVALS_PATH = LOG_DIR / "approvals.jsonl"

write_gate/policy.py

@@ -0,0 +1,49 @@
+"""Backward-compatible evaluate() entry point. Engine + guards do the real work."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from write_gate.catalog import Catalog
+from write_gate.config import Policy, demo_policy
+from write_gate.decision import (
+    RULE_EXPIRED,
+    RULE_OK,
+    RULE_PII,
+    RULE_SCHEMA,
+    Decision,
+    Evidence,
+)
+from write_gate.engine import evaluate as engine_evaluate
+
+__all__ = [
+    "Evidence",
+    "Decision",
+    "evaluate",
+    "RULE_OK",
+    "RULE_PII",
+    "RULE_EXPIRED",
+    "RULE_SCHEMA",
+]
+
+
+def evaluate(
+    sql: str,
+    catalog: Catalog,
+    policy: Policy | None = None,
+    conn: Any | None = None,
+    dialect: str = "duckdb",
+) -> Decision:
+    """Return a gate verdict for a single SQL string.
+
+    Library callers (tests / demo helpers) default to the demo policy so
+    existing legal-INSERT cases stay ALLOW. The CLI loads production
+    policy.yaml instead.
+    """
+    return engine_evaluate(
+        sql,
+        catalog,
+        policy=policy or demo_policy(),
+        conn=conn,
+        dialect=dialect,
+    )

write_gate/proxy.py

@@ -0,0 +1,273 @@
+"""SQL proxy in front of a real DB. Incoming SQL goes through WriteGate first.
+
+ALLOW then execute. BLOCK and REQUIRE_APPROVAL do not write. Uses
+WriteGate.execute only (never raw DuckDB). DuckDB is fully testable without
+a Postgres server; postgres:// / postgresql:// URLs take the same path.
+"""
+
+from __future__ import annotations
+
+import json
+import socket
+import sys
+from typing import Any, Callable, TextIO
+
+from write_gate.decision import ACTION_ALLOW, ACTION_APPROVAL, ACTION_BLOCK, Decision
+from write_gate.wrapper import WriteGate
+
+__all__ = [
+    "handle_sql",
+    "format_proxy",
+    "exit_code",
+    "split_statements",
+    "parse_listen_addr",
+    "serve_listen",
+    "handle_connection",
+    "run_cli",
+]
+
+
+def handle_sql(gate: WriteGate, sql: str) -> tuple[Decision, Any]:
+    """Gate one statement via WriteGate.execute (never raw DuckDB).
+
+    ALLOW executes. BLOCK and REQUIRE_APPROVAL return ``result=None``.
+    """
+    return gate.execute(sql)
+
+
+def exit_code(decision: Decision) -> int:
+    """Same as ``check``: 0 ALLOW, 1 REQUIRE_APPROVAL, 2 BLOCK."""
+    if decision.action == ACTION_ALLOW:
+        return 0
+    if decision.action == ACTION_APPROVAL:
+        return 1
+    return 2
+
+
+def _headline(action: str) -> str:
+    if action == ACTION_BLOCK:
+        return "BLOCKED"
+    if action == ACTION_APPROVAL:
+        return "APPROVAL REQUIRED"
+    return "ALLOWED"
+
+
+def _safe(value: object) -> str:
+    if value is None or value == "":
+        return "-"
+    return str(value)
+
+
+def _executed(decision: Decision, result: Any) -> bool:
+    return decision.action == ACTION_ALLOW and result is not None
+
+
+def format_proxy(decision: Decision, result: Any = None, *, as_json: bool = False) -> str:
+    """Human or JSON text for one proxy decision. Always includes executed."""
+    executed = _executed(decision, result)
+    if as_json:
+        payload = decision.to_dict()
+        payload["executed"] = executed
+        if executed:
+            try:
+                rows = result.fetchall()
+                payload["rows"] = [list(r) for r in rows]
+                payload["rowcount"] = len(payload["rows"])
+            except Exception:
+                payload["rowcount"] = getattr(result, "rowcount", None)
+        return json.dumps(payload, ensure_ascii=False, indent=2)
+    lines = [
+        _headline(decision.action),
+        f"Risk: {_safe(decision.risk)}",
+        f"Operation: {_safe(decision.operation).upper()}",
+        f"Table: {_safe(decision.table)}",
+        f"Rule: {_safe(decision.rule_id)}",
+        f"Reason: {_safe(decision.reason)}",
+        f"executed: {'yes' if executed else 'no'}",
+    ]
+    if decision.approval_id:
+        lines.append(f"Approval id: {decision.approval_id}")
+    return "\n".join(lines)
+
+
+def split_statements(text: str, *, once: bool = False) -> list[str]:
+    """Split stdin into SQL statements: one SQL, or one per non-empty line."""
+    raw = (text or "").strip()
+    if not raw:
+        return []
+    if once:
+        return [_strip_sql(raw)]
+    lines = [_strip_sql(ln) for ln in raw.splitlines()]
+    return [ln for ln in lines if ln]
+
+
+def _strip_sql(sql: str) -> str:
+    return sql.strip().rstrip(";").strip()
+
+
+def parse_listen_addr(spec: str) -> tuple[str, int]:
+    """Parse ``HOST:PORT`` (port 0 allowed for tests). Bare PORT binds 127.0.0.1."""
+    text = (spec or "").strip()
+    if not text:
+        raise ValueError("listen address is empty")
+    if ":" not in text:
+        return "127.0.0.1", int(text)
+    host, _, port_s = text.rpartition(":")
+    return (host or "127.0.0.1"), int(port_s)
+
+
+def read_sql_from_socket(sock: socket.socket, *, limit: int = 1_000_000) -> str:
+    """Read one SQL until newline or semicolon (or EOF / size limit)."""
+    buf = bytearray()
+    while len(buf) < limit:
+        chunk = sock.recv(4096)
+        if not chunk:
+            break
+        buf.extend(chunk)
+        if b"\n" in chunk or b";" in chunk:
+            break
+    text = bytes(buf).decode("utf-8", errors="replace")
+    for sep in ("\n", ";"):
+        if sep in text:
+            text = text.split(sep, 1)[0]
+            break
+    return _strip_sql(text)
+
+
+def handle_connection(
+    conn: socket.socket,
+    gate: WriteGate,
+    *,
+    as_json: bool = False,
+) -> Decision:
+    """One connection: read one SQL, gate, write BLOCKED/ALLOWED, close is caller."""
+    sql = read_sql_from_socket(conn)
+    if not sql:
+        body = "BLOCKED\nRule: empty_sql\nReason: empty SQL\nexecuted: no\n"
+        conn.sendall(body.encode("utf-8"))
+        return Decision(
+            action=ACTION_BLOCK,
+            risk="low",
+            rule_id="empty_sql",
+            reason="empty SQL",
+        )
+    decision, result = handle_sql(gate, sql)
+    body = format_proxy(decision, result, as_json=as_json) + "\n"
+    conn.sendall(body.encode("utf-8"))
+    return decision
+
+
+def serve_listen(
+    host: str,
+    port: int,
+    gate: WriteGate,
+    *,
+    as_json: bool = False,
+    stop: Callable[[], bool] | None = None,
+    on_bound: Callable[[str, int], None] | None = None,
+    err: TextIO | None = None,
+) -> int:
+    """Accept connections on HOST:PORT. Each connection is one SQL then close.
+
+    ``port=0`` lets the OS pick an ephemeral port (tests bind 127.0.0.1:0).
+    ``stop`` is polled between accepts so tests can shut the server down.
+    """
+    err = err if err is not None else sys.stderr
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    sock.bind((host, port))
+    sock.listen(8)
+    sock.settimeout(0.3)
+    bound_host, bound_port = sock.getsockname()[:2]
+    err.write(f"listening on {bound_host}:{bound_port}\n")
+    err.flush()
+    if on_bound is not None:
+        on_bound(bound_host, bound_port)
+    try:
+        while True:
+            if stop is not None and stop():
+                return 0
+            try:
+                conn, _addr = sock.accept()
+            except TimeoutError:
+                continue
+            except OSError:
+                if stop is not None and stop():
+                    return 0
+                raise
+            with conn:
+                handle_connection(conn, gate, as_json=as_json)
+    except KeyboardInterrupt:
+        err.write("proxy: listen stopped\n")
+        return 0
+    finally:
+        sock.close()
+
+
+def _worst_code(codes: list[int]) -> int:
+    if not codes:
+        return 2
+    return max(codes)
+
+
+def run_statements(
+    gate: WriteGate,
+    statements: list[str],
+    *,
+    as_json: bool = False,
+    out: TextIO | None = None,
+) -> int:
+    """Run each SQL through handle_sql. Exit code is the most severe result."""
+    out = out if out is not None else sys.stdout
+    codes: list[int] = []
+    for sql in statements:
+        decision, result = handle_sql(gate, sql)
+        out.write(format_proxy(decision, result, as_json=as_json) + "\n")
+        out.flush()
+        codes.append(exit_code(decision))
+    return _worst_code(codes)
+
+
+def run_cli(
+    gate: WriteGate,
+    *,
+    sql: str | None = None,
+    listen: str | None = None,
+    once: bool = False,
+    as_json: bool = False,
+    stdin: TextIO | None = None,
+    out: TextIO | None = None,
+    err: TextIO | None = None,
+    stop: Callable[[], bool] | None = None,
+    on_bound: Callable[[str, int], None] | None = None,
+) -> int:
+    """CLI body for ``sql-write-gate proxy``. One-shot --sql / stdin EOF / listen."""
+    err = err if err is not None else sys.stderr
+    stdin = stdin if stdin is not None else sys.stdin
+    if listen:
+        host, port = parse_listen_addr(listen)
+        return serve_listen(
+            host,
+            port,
+            gate,
+            as_json=as_json,
+            stop=stop,
+            on_bound=on_bound,
+            err=err,
+        )
+
+    once = bool(once) or bool(sql and str(sql).strip())
+    statements: list[str] = []
+    if sql and str(sql).strip():
+        statements = [_strip_sql(str(sql))]
+    else:
+        isatty = getattr(stdin, "isatty", lambda: False)
+        if callable(isatty) and isatty():
+            err.write("proxy: provide --sql, pipe SQL on stdin, or --listen HOST:PORT\n")
+            return 2
+        statements = split_statements(stdin.read(), once=once)
+
+    if not statements:
+        err.write("proxy: empty SQL\n")
+        return 2
+    return run_statements(gate, statements, as_json=as_json, out=out)