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

write_gate/config.py

@@ -0,0 +1,117 @@
+"""Load policy.yaml (environment, per-operation rules, blast-radius limits)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from write_gate.paths import POLICY_PATH
+
+VALID_RULES = {"allow", "block", "approval"}
+VALID_OPS = ("select", "insert", "update", "delete", "ddl")
+
+PRODUCTION_DEFAULTS: dict[str, Any] = {
+    "environment": "production",
+    "rules": {
+        "select": "allow",
+        "insert": "approval",
+        "update": "approval",
+        "delete": "block",
+        "ddl": "block",
+    },
+    "limits": {
+        "update_rows": 100,
+        "delete_rows": 50,
+    },
+}
+
+DEMO_DEFAULTS: dict[str, Any] = {
+    "environment": "demo",
+    "rules": {
+        "select": "allow",
+        "insert": "allow",
+        "update": "allow",
+        "delete": "allow",
+        "ddl": "block",
+    },
+    "limits": {
+        "update_rows": 10000,
+        "delete_rows": 10000,
+    },
+}
+
+
+@dataclass(frozen=True)
+class Policy:
+    environment: str
+    rules: dict[str, str] = field(default_factory=dict)
+    update_rows: int = 100
+    delete_rows: int = 50
+
+    def rule_for(self, operation: str) -> str:
+        op = (operation or "ddl").lower()
+        return self.rules.get(op, "block")
+
+    def row_limit(self, operation: str) -> int | None:
+        if operation == "update":
+            return self.update_rows
+        if operation == "delete":
+            return self.delete_rows
+        return None
+
+    def with_env_approvals_cleared(self) -> "Policy":
+        """Human approve clears environment 'approval' rules only (block stays)."""
+        rules = {
+            op: ("allow" if rule == "approval" else rule)
+            for op, rule in self.rules.items()
+        }
+        return Policy(
+            environment=self.environment,
+            rules=rules,
+            update_rows=self.update_rows,
+            delete_rows=self.delete_rows,
+        )
+
+
+def _normalize_rules(raw: Any) -> dict[str, str]:
+    src = dict(PRODUCTION_DEFAULTS["rules"])
+    if isinstance(raw, dict):
+        for key, value in raw.items():
+            k = str(key).lower()
+            v = str(value).lower()
+            if k in VALID_OPS and v in VALID_RULES:
+                src[k] = v
+    return {k: src[k] for k in VALID_OPS}
+
+
+def policy_from_dict(raw: dict[str, Any] | None = None) -> Policy:
+    data = raw or {}
+    limits = data.get("limits") or {}
+    return Policy(
+        environment=str(data.get("environment") or PRODUCTION_DEFAULTS["environment"]),
+        rules=_normalize_rules(data.get("rules")),
+        update_rows=int(limits.get("update_rows", PRODUCTION_DEFAULTS["limits"]["update_rows"])),
+        delete_rows=int(limits.get("delete_rows", PRODUCTION_DEFAULTS["limits"]["delete_rows"])),
+    )
+
+
+def load_policy(path: Path | str | None = None) -> Policy:
+    policy_path = Path(path) if path else POLICY_PATH
+    if not policy_path.exists():
+        return policy_from_dict(PRODUCTION_DEFAULTS)
+    with policy_path.open(encoding="utf-8") as fh:
+        raw = yaml.safe_load(fh) or {}
+    if not isinstance(raw, dict):
+        return policy_from_dict(PRODUCTION_DEFAULTS)
+    return policy_from_dict(raw)
+
+
+def production_policy() -> Policy:
+    return policy_from_dict(PRODUCTION_DEFAULTS)
+
+
+def demo_policy() -> Policy:
+    return policy_from_dict(DEMO_DEFAULTS)

write_gate/db.py

@@ -0,0 +1,5 @@
+"""Compatibility shim. Prefer write_gate.adapters.duckdb."""
+
+from write_gate.adapters.duckdb import ORDERS_DDL, connect, execute_user_sql
+
+__all__ = ["ORDERS_DDL", "connect", "execute_user_sql"]

write_gate/decision.py

@@ -0,0 +1,143 @@
+"""Unified decision model for the write gate."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+ACTION_ALLOW = "ALLOW"
+ACTION_BLOCK = "BLOCK"
+ACTION_APPROVAL = "REQUIRE_APPROVAL"
+
+RISK_LOW = "low"
+RISK_MEDIUM = "medium"
+RISK_CRITICAL = "critical"
+
+VERDICT_PASS = "PASS"
+VERDICT_WARN = "WARN"
+VERDICT_APPROVAL = "APPROVAL"
+VERDICT_BLOCK = "BLOCK"
+
+RULE_OK = "ok"
+RULE_PII = "pii_column"
+RULE_RESTRICTED = "restricted_column"
+RULE_EXPIRED = "expired_partition"
+RULE_SCHEMA = "schema_mismatch"
+RULE_DELETE_NO_WHERE = "delete_without_where"
+RULE_UPDATE_NO_WHERE = "update_without_where"
+RULE_DROP_TABLE = "drop_table"
+RULE_TRUNCATE = "truncate_table"
+RULE_ALTER_TABLE = "alter_table"
+RULE_BLAST = "blast_radius_exceeded"
+RULE_ENV = "environment_policy"
+RULE_RAW_DB_CLI = "raw_db_cli"
+
+
+@dataclass
+class GuardResult:
+    """One guard's verdict. Engine reduces a list of these to a Decision."""
+
+    name: str
+    verdict: str = VERDICT_PASS
+    rule_id: str | None = None
+    reason: str = ""
+    risk: str = RISK_LOW
+    evidence: dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def pass_(cls, name: str, **kwargs: Any) -> "GuardResult":
+        return cls(name=name, verdict=VERDICT_PASS, **kwargs)
+
+    @classmethod
+    def warn(cls, name: str, rule_id: str, reason: str, **kwargs: Any) -> "GuardResult":
+        return cls(
+            name=name,
+            verdict=VERDICT_WARN,
+            rule_id=rule_id,
+            reason=reason,
+            risk=kwargs.pop("risk", RISK_MEDIUM),
+            **kwargs,
+        )
+
+    @classmethod
+    def approval(
+        cls,
+        name: str,
+        rule_id: str,
+        reason: str,
+        *,
+        risk: str = RISK_MEDIUM,
+        evidence: dict[str, Any] | None = None,
+    ) -> "GuardResult":
+        return cls(
+            name=name,
+            verdict=VERDICT_APPROVAL,
+            rule_id=rule_id,
+            reason=reason,
+            risk=risk,
+            evidence=evidence or {},
+        )
+
+    @classmethod
+    def block(
+        cls,
+        name: str,
+        rule_id: str,
+        reason: str,
+        *,
+        risk: str = RISK_CRITICAL,
+        evidence: dict[str, Any] | None = None,
+    ) -> "GuardResult":
+        return cls(
+            name=name,
+            verdict=VERDICT_BLOCK,
+            rule_id=rule_id,
+            reason=reason,
+            risk=risk,
+            evidence=evidence or {},
+        )
+
+
+@dataclass
+class Decision:
+    """Final engine decision. `allowed` stays True only for ALLOW (execute path)."""
+
+    action: str
+    risk: str
+    rule_id: str
+    reason: str
+    evidence: dict[str, Any] = field(default_factory=dict)
+    sql: str = ""
+    operation: str | None = None
+    table: str | None = None
+    estimated_rows: int | None = None
+    approval_id: str | None = None
+
+    @property
+    def allowed(self) -> bool:
+        return self.action == ACTION_ALLOW
+
+    @property
+    def message(self) -> str:
+        return self.reason
+
+    def to_dict(self) -> dict[str, Any]:
+        payload = {
+            "allowed": self.allowed,
+            "rule_id": self.rule_id,
+            "message": self.reason,
+            "sql": self.sql,
+            "action": self.action,
+            "risk": self.risk,
+            "reason": self.reason,
+            "evidence": self.evidence,
+            "operation": self.operation,
+            "table": self.table,
+            "estimated_rows": self.estimated_rows,
+            "approval_id": self.approval_id,
+        }
+        return payload
+
+
+# Backward-compatible alias used by existing tests / demo.
+Evidence = Decision

write_gate/engine.py

@@ -0,0 +1,150 @@
+"""Policy engine: run guards, reduce to a single Decision.
+
+any BLOCK → BLOCK; else any APPROVAL → REQUIRE_APPROVAL; else ALLOW.
+Guard order prefers specific dangerous-SQL rules over environment policy.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+from write_gate.catalog import Catalog
+from write_gate.config import Policy, production_policy
+from write_gate.decision import (
+    ACTION_ALLOW,
+    ACTION_APPROVAL,
+    ACTION_BLOCK,
+    RISK_LOW,
+    RULE_OK,
+    VERDICT_APPROVAL,
+    VERDICT_BLOCK,
+    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
+
+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:
+    sql: str
+    parsed: ParsedSQL
+    catalog: Catalog
+    policy: Policy
+    conn: Any | None = None
+    dialect: str = BACKEND_DUCKDB
+    guard_results: list[GuardResult] = field(default_factory=list)
+
+
+def evaluate(
+    sql: str,
+    catalog: Catalog,
+    policy: Policy | None = None,
+    conn: Any | None = None,
+    dialect: str = BACKEND_DUCKDB,
+) -> Decision:
+    parsed = parse(sql, dialect=sqlglot_dialect(dialect))
+    ctx = Context(
+        sql=sql,
+        parsed=parsed,
+        catalog=catalog,
+        policy=policy or production_policy(),
+        conn=conn,
+        dialect=dialect,
+    )
+    results = [guard(ctx) for guard in GUARDS]
+    ctx.guard_results = results
+    return reduce(ctx, results)
+
+
+def reduce(ctx: Context, results: list[GuardResult]) -> Decision:
+    parsed = ctx.parsed
+    estimated = _first_estimated(results)
+    evidence_acc: dict[str, Any] = {}
+    for result in results:
+        if result.evidence:
+            evidence_acc[result.name] = result.evidence
+
+    blocks = [r for r in results if r.verdict == VERDICT_BLOCK]
+    approvals = [r for r in results if r.verdict == VERDICT_APPROVAL]
+
+    if blocks:
+        chosen = blocks[0]
+        return Decision(
+            action=ACTION_BLOCK,
+            risk=chosen.risk,
+            rule_id=chosen.rule_id or RULE_OK,
+            reason=chosen.reason,
+            evidence={**evidence_acc, **chosen.evidence},
+            sql=ctx.sql,
+            operation=parsed.operation,
+            table=parsed.table,
+            estimated_rows=estimated if estimated is not None else chosen.evidence.get("estimated_rows"),
+        )
+    if approvals:
+        chosen = approvals[0]
+        return Decision(
+            action=ACTION_APPROVAL,
+            risk=chosen.risk,
+            rule_id=chosen.rule_id or RULE_OK,
+            reason=chosen.reason,
+            evidence={**evidence_acc, **chosen.evidence},
+            sql=ctx.sql,
+            operation=parsed.operation,
+            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,
+        operation=parsed.operation,
+        table=parsed.table,
+        estimated_rows=estimated,
+    )
+
+
+def _first_estimated(results: list[GuardResult]) -> int | None:
+    for result in results:
+        value = result.evidence.get("estimated_rows") if result.evidence else None
+        if isinstance(value, int):
+            return value
+    return None
+
+
+def _allow_reason(parsed: ParsedSQL) -> str:
+    if parsed.operation == "select":
+        return "只读 SELECT，绕过写库门禁"
+    table = parsed.table or "?"
+    if parsed.operation == "insert":
+        return f"写入通过门禁: 表 {table}，列 {parsed.write_columns}，分区未过期且不含 PII"
+    if parsed.operation == "update":
+        return f"UPDATE 通过门禁: 表 {table}，列 {parsed.write_columns}"
+    if parsed.operation == "delete":
+        return f"DELETE 通过门禁: 表 {table}"
+    return f"{parsed.operation.upper()} 通过门禁: 表 {table}"

write_gate/guards/__init__.py

@@ -0,0 +1,17 @@
+"""Guard functions. Each returns PASS | WARN | APPROVAL | BLOCK."""
+
+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.freshness import check_freshness
+from write_gate.guards.pii import check_pii
+from write_gate.guards.schema import check_schema
+
+__all__ = [
+    "check_blast_radius",
+    "check_destructive",
+    "check_environment",
+    "check_freshness",
+    "check_pii",
+    "check_schema",
+]

write_gate/guards/blast_radius.py

@@ -0,0 +1,73 @@
+"""Blast-radius guard: estimate affected rows for UPDATE/DELETE before execute."""
+
+from __future__ import annotations
+
+from write_gate.adapters.base import BACKEND_DUCKDB, count_sql, sqlglot_dialect
+from write_gate.decision import RULE_BLAST, RISK_CRITICAL, GuardResult
+from write_gate.parser import where_sql
+
+NAME = "blast_radius"
+
+
+def check_blast_radius(ctx) -> GuardResult:
+    parsed = ctx.parsed
+    if parsed.statement is None or parsed.error:
+        return GuardResult.pass_(NAME)
+    if parsed.operation not in {"update", "delete"}:
+        return GuardResult.pass_(NAME)
+
+    table = parsed.table
+    if not table:
+        return GuardResult.pass_(NAME)
+
+    limit = ctx.policy.row_limit(parsed.operation)
+    if limit is None:
+        return GuardResult.pass_(NAME)
+
+    conn = ctx.conn
+    if conn is None:
+        return GuardResult.pass_(
+            NAME,
+            evidence={"skipped": True, "reason": "no connection to estimate rows"},
+        )
+
+    dialect = getattr(ctx, "dialect", BACKEND_DUCKDB)
+    estimated = _estimate_rows(conn, table, parsed.where, dialect=dialect)
+    evidence = {
+        "estimated_rows": estimated,
+        "max_rows": limit,
+        "operation": parsed.operation,
+        "table": table,
+    }
+    if estimated is None:
+        return GuardResult.pass_(NAME, evidence={**evidence, "skipped": True})
+    if estimated > limit:
+        return GuardResult.block(
+            NAME,
+            RULE_BLAST,
+            (
+                f"{parsed.operation.upper()} on {table} would affect {estimated} rows, "
+                f"exceeding max {limit}"
+            ),
+            risk=RISK_CRITICAL,
+            evidence=evidence,
+        )
+    return GuardResult.pass_(NAME, evidence=evidence)
+
+
+def _estimate_rows(conn, table: str, where, dialect: str = BACKEND_DUCKDB) -> int | None:
+    # Table name comes from the parser identifier, not raw user interpolation of extra SQL.
+    if not table.isidentifier():
+        return None
+    predicate = where_sql(where, dialect=sqlglot_dialect(dialect))
+    sql = count_sql(table, predicate, backend=dialect)
+    try:
+        row = conn.execute(sql).fetchone()
+    except Exception:
+        return None
+    if not row:
+        return 0
+    try:
+        return int(row[0])
+    except (TypeError, ValueError):
+        return None

write_gate/guards/destructive.py

@@ -0,0 +1,85 @@
+"""Dangerous SQL guard: unbounded DELETE/UPDATE, DROP, TRUNCATE, ALTER."""
+
+from __future__ import annotations
+
+from sqlglot import exp
+
+from write_gate.decision import (
+    RULE_ALTER_TABLE,
+    RULE_DELETE_NO_WHERE,
+    RULE_DROP_TABLE,
+    RULE_TRUNCATE,
+    RULE_UPDATE_NO_WHERE,
+    GuardResult,
+)
+
+NAME = "destructive"
+
+
+def check_destructive(ctx) -> GuardResult:
+    parsed = ctx.parsed
+    stmt = parsed.statement
+    if stmt is None:
+        return GuardResult.pass_(NAME)
+
+    if isinstance(stmt, exp.Delete) and not parsed.has_where:
+        table = parsed.table or "?"
+        return GuardResult.block(
+            NAME,
+            RULE_DELETE_NO_WHERE,
+            f"DELETE without a WHERE clause is forbidden (full-table delete on {table})",
+            evidence={"table": table, "has_where": False},
+        )
+
+    if isinstance(stmt, exp.Update) and not parsed.has_where:
+        table = parsed.table or "?"
+        return GuardResult.block(
+            NAME,
+            RULE_UPDATE_NO_WHERE,
+            f"UPDATE without a WHERE clause is forbidden (full-table update on {table})",
+            evidence={"table": table, "has_where": False},
+        )
+
+    if isinstance(stmt, exp.Drop):
+        kind = (stmt.args.get("kind") or "TABLE")
+        kind_s = str(kind).upper()
+        table = parsed.table or ident_fallback(stmt)
+        if kind_s in {"TABLE", "VIEW", ""} or kind is None:
+            return GuardResult.block(
+                NAME,
+                RULE_DROP_TABLE,
+                f"DROP {kind_s or 'TABLE'} {table} is forbidden",
+                evidence={"table": table, "kind": kind_s},
+            )
+        return GuardResult.block(
+            NAME,
+            RULE_DROP_TABLE,
+            f"DROP {kind_s} is forbidden",
+            evidence={"kind": kind_s},
+        )
+
+    if isinstance(stmt, exp.TruncateTable) or type(stmt).__name__ == "TruncateTable":
+        table = parsed.table or "?"
+        return GuardResult.block(
+            NAME,
+            RULE_TRUNCATE,
+            f"TRUNCATE TABLE {table} is forbidden",
+            evidence={"table": table},
+        )
+
+    if isinstance(stmt, exp.Alter) or type(stmt).__name__ in {"Alter", "AlterTable"}:
+        table = parsed.table or "?"
+        return GuardResult.block(
+            NAME,
+            RULE_ALTER_TABLE,
+            f"ALTER TABLE {table} is forbidden",
+            evidence={"table": table},
+        )
+
+    return GuardResult.pass_(NAME)
+
+
+def ident_fallback(stmt) -> str:
+    this = getattr(stmt, "this", None)
+    name = getattr(this, "name", None) if this is not None else None
+    return str(name).lower() if name else "?"

write_gate/guards/environment.py

@@ -0,0 +1,41 @@
+"""Environment policy: per-operation allow | block | approval."""
+
+from __future__ import annotations
+
+from write_gate.decision import RULE_ENV, RISK_CRITICAL, RISK_MEDIUM, GuardResult
+
+NAME = "environment"
+
+
+def check_environment(ctx) -> GuardResult:
+    parsed = ctx.parsed
+    if parsed.statement is None or parsed.error:
+        return GuardResult.pass_(NAME)
+
+    operation = parsed.operation if parsed.operation != "unknown" else "ddl"
+    rule = ctx.policy.rule_for(operation)
+    env = ctx.policy.environment
+    evidence = {
+        "environment": env,
+        "operation": operation,
+        "policy_rule": rule,
+    }
+    if rule == "allow":
+        return GuardResult.pass_(NAME, evidence=evidence)
+    if rule == "approval":
+        return GuardResult.approval(
+            NAME,
+            RULE_ENV,
+            f"{operation.upper()} requires approval in {env}",
+            risk=RISK_MEDIUM,
+            evidence=evidence,
+        )
+    # block
+    risk = RISK_CRITICAL if operation in {"delete", "ddl", "update"} else RISK_MEDIUM
+    return GuardResult.block(
+        NAME,
+        RULE_ENV,
+        f"{operation.upper()} is blocked by policy in {env}",
+        risk=risk,
+        evidence=evidence,
+    )