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

write_gate/guards/freshness.py

@@ -0,0 +1,102 @@
+"""Freshness guard: expired partitions and stale/non-writable tables."""
+
+from __future__ import annotations
+
+from datetime import date
+
+from write_gate.decision import RULE_EXPIRED, RULE_SCHEMA, RISK_MEDIUM, GuardResult
+from write_gate.parser import partition_dates_from_insert, partition_dates_from_where
+
+NAME = "freshness"
+
+
+def check_freshness(ctx) -> GuardResult:
+    parsed = ctx.parsed
+    if parsed.statement is None or parsed.error:
+        return GuardResult.pass_(NAME)
+    if parsed.operation not in {"insert", "update", "delete"}:
+        return GuardResult.pass_(NAME)
+
+    table_name = parsed.table
+    if not table_name:
+        return GuardResult.pass_(NAME)
+    spec = ctx.catalog.table(table_name)
+    if spec is None:
+        return GuardResult.pass_(NAME)
+
+    if spec.stale or not spec.writable:
+        return GuardResult.block(
+            NAME,
+            RULE_EXPIRED,
+            f"表 {spec.name} 已标记为 stale/不可写，拒绝全部写入",
+            risk=RISK_MEDIUM,
+            evidence={"table": spec.name, "stale": spec.stale, "writable": spec.writable},
+        )
+
+    part = spec.partition_column
+    if not part:
+        return GuardResult.pass_(NAME)
+
+    if parsed.operation == "insert":
+        rows = parsed.insert_rows or []
+        cols = parsed.write_columns
+        dates = partition_dates_from_insert(cols, rows, part)
+        return _evaluate_dates(
+            spec,
+            ctx.catalog,
+            dates,
+            missing_ok=False,
+            missing_message=f"INSERT 必须显式写出分区列 {part}",
+        )
+
+    dates = partition_dates_from_where(parsed.where, part)
+    # UPDATE/DELETE: if WHERE names an expired partition, block.
+    # If WHERE exists but does not mention the partition, let blast_radius handle scope.
+    return _evaluate_dates(
+        spec,
+        ctx.catalog,
+        dates,
+        missing_ok=True,
+        missing_message=(
+            f"{parsed.operation.upper()} 必须在 WHERE 中约束分区列 {part}，避免误写过期分区"
+        ),
+    )
+
+
+def _evaluate_dates(
+    spec,
+    catalog,
+    partition_dates: list[date | None],
+    *,
+    missing_ok: bool,
+    missing_message: str,
+) -> GuardResult:
+    present = [d for d in partition_dates if d is not None]
+    if not present and not missing_ok:
+        return GuardResult.block(
+            NAME,
+            RULE_SCHEMA,
+            missing_message,
+            risk=RISK_MEDIUM,
+            evidence={"partition_column": spec.partition_column},
+        )
+    cutoff = catalog.cutoff_date
+    expired = [d for d in present if d < cutoff]
+    if expired:
+        worst = min(expired)
+        return GuardResult.block(
+            NAME,
+            RULE_EXPIRED,
+            (
+                f"分区 {spec.partition_column}={worst.isoformat()} 早于新鲜度截止日期 "
+                f"{cutoff.isoformat()}（as_of={catalog.as_of_date.isoformat()}, "
+                f"freshness_days={catalog.freshness_days}），拒绝写入"
+            ),
+            risk=RISK_MEDIUM,
+            evidence={
+                "partition_column": spec.partition_column,
+                "partition_value": worst.isoformat(),
+                "cutoff": cutoff.isoformat(),
+            },
+        )
+    return GuardResult.pass_(NAME)

write_gate/guards/pii.py

@@ -0,0 +1,84 @@
+"""PII / restricted-column guard.
+
+Writes to PII or restricted columns are BLOCK.
+SELECT of PII columns is REQUIRE_APPROVAL (not silent allow).
+SELECT of restricted columns is BLOCK.
+"""
+
+from __future__ import annotations
+
+from write_gate.decision import RULE_PII, RULE_RESTRICTED, RISK_CRITICAL, RISK_MEDIUM, GuardResult
+
+NAME = "pii"
+
+
+def check_pii(ctx) -> GuardResult:
+    parsed = ctx.parsed
+    if parsed.statement is None or parsed.error:
+        return GuardResult.pass_(NAME)
+
+    table_name = parsed.table
+    if not table_name:
+        return GuardResult.pass_(NAME)
+    spec = ctx.catalog.table(table_name)
+    if spec is None:
+        return GuardResult.pass_(NAME)
+
+    if parsed.operation in {"insert", "update"}:
+        cols = list(parsed.write_columns)
+        restricted_hit = [c for c in cols if c in spec.restricted_columns]
+        if restricted_hit:
+            return GuardResult.block(
+                NAME,
+                RULE_RESTRICTED,
+                (
+                    f"禁止写入受限列 {restricted_hit}"
+                    f"（表 {spec.name} 的 restricted_columns={sorted(spec.restricted_columns)}）"
+                ),
+                risk=RISK_CRITICAL,
+                evidence={"columns": restricted_hit, "table": spec.name},
+            )
+        pii_hit = [c for c in cols if c in spec.pii_columns]
+        if pii_hit:
+            return GuardResult.block(
+                NAME,
+                RULE_PII,
+                (
+                    f"禁止写入 PII 列 {pii_hit}"
+                    f"（表 {spec.name} 的 pii_columns={sorted(spec.pii_columns)}）"
+                ),
+                risk=RISK_CRITICAL,
+                evidence={"columns": pii_hit, "table": spec.name},
+            )
+        return GuardResult.pass_(NAME)
+
+    if parsed.operation == "select":
+        selected = _selected_columns(parsed, spec)
+        restricted_hit = [c for c in selected if c in spec.restricted_columns]
+        if restricted_hit:
+            return GuardResult.block(
+                NAME,
+                RULE_RESTRICTED,
+                f"禁止 SELECT 受限列 {restricted_hit}（表 {spec.name}）",
+                risk=RISK_CRITICAL,
+                evidence={"columns": restricted_hit, "table": spec.name},
+            )
+        pii_hit = [c for c in selected if c in spec.pii_columns]
+        if pii_hit:
+            return GuardResult.approval(
+                NAME,
+                RULE_PII,
+                (
+                    f"SELECT 包含 PII 列 {pii_hit}，需要审批"
+                    f"（表 {spec.name} 的 pii_columns={sorted(spec.pii_columns)}）"
+                ),
+                risk=RISK_MEDIUM,
+                evidence={"columns": pii_hit, "table": spec.name},
+            )
+    return GuardResult.pass_(NAME)
+
+
+def _selected_columns(parsed, spec) -> list[str]:
+    if parsed.star:
+        return list(spec.columns.keys())
+    return list(parsed.select_columns or parsed.columns)

write_gate/guards/schema.py

@@ -0,0 +1,147 @@
+"""Schema guard: parse errors, unknown table/column, type mismatch, unsupported SQL."""
+
+from __future__ import annotations
+
+from sqlglot import exp
+
+from write_gate.catalog import TableSpec
+from write_gate.decision import RULE_SCHEMA, GuardResult
+from write_gate.parser import literal_value, type_ok
+
+NAME = "schema"
+
+
+def check_schema(ctx) -> GuardResult:
+    parsed = ctx.parsed
+    if parsed.error:
+        return GuardResult.block(
+            NAME,
+            parsed.error_rule or RULE_SCHEMA,
+            parsed.error,
+            risk="medium",
+        )
+
+    stmt = parsed.statement
+    if stmt is None:
+        return GuardResult.block(NAME, RULE_SCHEMA, "SQL 无法解析", risk="medium")
+
+    operation = parsed.operation
+    if operation == "select":
+        return GuardResult.pass_(NAME)
+
+    if operation == "ddl":
+        # Destructive/environment guards own DROP/ALTER/TRUNCATE/CREATE.
+        return GuardResult.pass_(NAME)
+
+    if operation not in {"insert", "update", "delete"}:
+        return GuardResult.block(
+            NAME,
+            RULE_SCHEMA,
+            f"不支持的写语句类型 {type(stmt).__name__}，仅允许 INSERT/UPDATE/DELETE",
+            risk="medium",
+        )
+
+    table_name = parsed.table
+    if not table_name:
+        return GuardResult.block(NAME, RULE_SCHEMA, "无法从 SQL 中解析目标表", risk="medium")
+
+    spec = ctx.catalog.table(table_name)
+    if spec is None:
+        return GuardResult.block(
+            NAME,
+            RULE_SCHEMA,
+            f"未知表 {table_name}，不在目录中",
+            risk="medium",
+            evidence={"table": table_name},
+        )
+
+    if isinstance(stmt, exp.Insert):
+        return _check_insert(parsed, spec)
+    if isinstance(stmt, exp.Update):
+        return _check_update(parsed, spec)
+    return GuardResult.pass_(NAME)
+
+
+def _check_insert(parsed, spec: TableSpec) -> GuardResult:
+    cols = list(parsed.write_columns)
+    if not cols:
+        cols = list(spec.columns.keys())
+        parsed.write_columns = cols
+        parsed.columns = cols
+    if any(c == "" for c in cols):
+        return GuardResult.block(NAME, RULE_SCHEMA, "INSERT 列名无法解析", risk="medium")
+
+    rows = parsed.insert_rows
+    if rows is None:
+        return GuardResult.block(
+            NAME,
+            RULE_SCHEMA,
+            "仅支持 INSERT ... VALUES (...); INSERT ... SELECT 未开放",
+            risk="medium",
+        )
+
+    for row in rows:
+        if len(row) != len(cols):
+            return GuardResult.block(
+                NAME,
+                RULE_SCHEMA,
+                f"INSERT 列数 {len(cols)} 与值个数 {len(row)} 不一致",
+                risk="medium",
+            )
+        assignments = dict(zip(cols, row))
+        failed = _columns_and_types(spec, cols, assignments)
+        if failed:
+            return failed
+    return GuardResult.pass_(NAME)
+
+
+def _check_update(parsed, spec: TableSpec) -> GuardResult:
+    assignments = parsed.assignments
+    cols = list(assignments.keys())
+    if not cols:
+        return GuardResult.block(NAME, RULE_SCHEMA, "UPDATE 未解析到 SET 列", risk="medium")
+    failed = _columns_and_types(spec, cols, assignments)
+    if failed:
+        return failed
+    return GuardResult.pass_(NAME)
+
+
+def _columns_and_types(
+    spec: TableSpec,
+    write_cols: list[str],
+    assignments: dict,
+) -> GuardResult | None:
+    unknown = [c for c in write_cols if c not in spec.columns]
+    if unknown:
+        return GuardResult.block(
+            NAME,
+            RULE_SCHEMA,
+            f"未知列 {unknown}，表 {spec.name} 的列为 {sorted(spec.columns)}",
+            risk="medium",
+            evidence={"unknown_columns": unknown, "table": spec.name},
+        )
+
+    for col in write_cols:
+        node = assignments.get(col)
+        expected = spec.columns[col]
+        if node is not None and not type_ok(expected, node):
+            got = literal_value(node)
+            return GuardResult.block(
+                NAME,
+                RULE_SCHEMA,
+                f"列 {col} 类型不匹配: 期望 {expected}，实际值 {got!r}",
+                risk="medium",
+                evidence={"column": col, "expected": expected, "value": got},
+            )
+
+    skip = spec.pii_columns | spec.restricted_columns
+    not_allowed = [c for c in write_cols if c not in spec.allowed_write_columns and c not in skip]
+    if not_allowed:
+        return GuardResult.block(
+            NAME,
+            RULE_SCHEMA,
+            f"列 {not_allowed} 不在允许写入列表 {sorted(spec.allowed_write_columns)}",
+            risk="medium",
+            evidence={"not_allowed": not_allowed},
+        )
+    return None

write_gate/hooks.py

@@ -0,0 +1,360 @@
+"""PreToolUse hook: agents cannot talk to the DB via raw psql / other DB CLIs.
+
+The hook never executes user SQL. Extracted statements go through WriteGate.check
+only. Interactive or unparseable DB shells are refused.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import shlex
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, TextIO
+
+from write_gate.decision import ACTION_ALLOW, ACTION_APPROVAL, RULE_RAW_DB_CLI, Decision
+from write_gate.wrapper import WriteGate
+
+DB_CLIS = ("psql", "mysql", "mysqlsh", "duckdb", "sqlite3")
+
+_SQL_FLAGS: dict[str, tuple[str, ...]] = {
+    "psql": ("-c", "--command"),
+    "mysql": ("-e", "--execute"),
+    "mysqlsh": ("-e", "--execute"),
+    "duckdb": ("-c", "--cmd", "-s", "--sql"),
+    "sqlite3": ("-cmd",),
+}
+
+_WRAPPERS = {"sudo", "env", "command", "nohup", "time", "nice", "stdbuf", "timeout"}
+
+_SHELL_OPS = {"&&", "||", "|", ";", "&"}
+
+_DB_CLI_WORD = re.compile(
+    r"(?:^|[^\w.-])(psql|mysqlsh|mysql|duckdb|sqlite3)(?:$|[^\w.-])",
+    re.IGNORECASE,
+)
+
+_ENV_ASSIGN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=")
+
+_SQL_HEAD = re.compile(
+    r"^(SELECT|INSERT|UPDATE|DELETE|DROP|TRUNCATE|ALTER|CREATE|WITH|COPY|"
+    r"REPLACE|MERGE|CALL|DO|BEGIN|EXPLAIN|VACUUM|ANALYZE|GRANT|REVOKE)\b",
+    re.IGNORECASE,
+)
+
+
+@dataclass(frozen=True)
+class _Segment:
+    cli: str | None
+    sqls: tuple[str, ...]
+    raw: bool
+
+
+def command_from_payload(payload: Any) -> str | None:
+    """Pull the bash command out of a Claude / Codex PreToolUse JSON object."""
+    if not isinstance(payload, dict):
+        return None
+    tool_input = payload.get("tool_input")
+    if isinstance(tool_input, dict):
+        for key in ("command", "cmd", "sql"):
+            value = tool_input.get(key)
+            if isinstance(value, str) and value.strip():
+                return value
+    elif isinstance(tool_input, str) and tool_input.strip():
+        return tool_input
+    for key in ("command", "cmd"):
+        value = payload.get(key)
+        if isinstance(value, str) and value.strip():
+            return value
+    return None
+
+
+def parse_stdin_payload(text: str) -> Any | None:
+    text = (text or "").strip()
+    if not text:
+        return None
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        return None
+
+
+def _basename(token: str) -> str:
+    name = Path(token).name.lower()
+    if name.endswith(".exe"):
+        name = name[:-4]
+    return name
+
+
+def _is_sql_flag(cli: str, token: str) -> bool:
+    return token in _SQL_FLAGS.get(cli, ())
+
+
+def _looks_like_sql(text: str) -> bool:
+    return bool(_SQL_HEAD.match(text.strip().lstrip("(")))
+
+
+def _skip_wrappers(tokens: list[str]) -> list[str]:
+    i = 0
+    n = len(tokens)
+    while i < n:
+        t = tokens[i]
+        if _ENV_ASSIGN.match(t):
+            i += 1
+            continue
+        name = _basename(t)
+        if name in _WRAPPERS:
+            i += 1
+            while i < n and tokens[i].startswith("-"):
+                i += 1
+            # timeout DURATION cmd
+            if name == "timeout" and i < n and not tokens[i].startswith("-"):
+                i += 1
+            continue
+        break
+    return tokens[i:]
+
+
+def _split_segments(tokens: list[str]) -> list[list[str]]:
+    segments: list[list[str]] = []
+    current: list[str] = []
+    for t in tokens:
+        if t in _SHELL_OPS:
+            if current:
+                segments.append(current)
+                current = []
+            continue
+        current.append(t)
+    if current:
+        segments.append(current)
+    return segments
+
+
+def _collect_sql_tokens(tokens: list[str], start: int) -> tuple[list[str], int]:
+    if start >= len(tokens) or tokens[start].startswith("-"):
+        return [], start
+    first = tokens[start]
+    if any(ch.isspace() for ch in first):
+        return [first], start + 1
+    parts: list[str] = []
+    i = start
+    while i < len(tokens) and not tokens[i].startswith("-"):
+        parts.append(tokens[i])
+        i += 1
+    return parts, i
+
+
+def _extract_from_cli(cli: str, tokens: list[str]) -> tuple[list[str], bool]:
+    """Return (sql_strings, saw_sql_flag_without_value)."""
+    sqls: list[str] = []
+    incomplete_flag = False
+    positionals: list[str] = []
+    i = 1
+    while i < len(tokens):
+        t = tokens[i]
+        if t.startswith("--") and "=" in t:
+            flag, _, value = t.partition("=")
+            if _is_sql_flag(cli, flag):
+                extra, i = _collect_sql_tokens(tokens, i + 1)
+                joined = " ".join([value, *extra]).strip() if extra else value.strip()
+                if joined:
+                    sqls.append(joined)
+                else:
+                    incomplete_flag = True
+                continue
+            i += 1
+            continue
+        if _is_sql_flag(cli, t):
+            extra, i = _collect_sql_tokens(tokens, i + 1)
+            if extra:
+                sqls.append(" ".join(extra).strip())
+            else:
+                incomplete_flag = True
+            continue
+        if t.startswith("-"):
+            i += 1
+            if i < len(tokens) and not tokens[i].startswith("-"):
+                i += 1
+            continue
+        positionals.append(t)
+        i += 1
+
+    if sqls:
+        return sqls, incomplete_flag
+    if incomplete_flag:
+        return [], True
+    if cli in {"duckdb", "sqlite3"}:
+        if len(positionals) >= 2:
+            return [positionals[1]], False
+        if len(positionals) == 1 and _looks_like_sql(positionals[0]):
+            return [positionals[0]], False
+    return [], False
+
+
+def inspect_bash(command: str) -> list[_Segment]:
+    """Inspect a bash command for DB CLIs. Empty list means not a DB CLI."""
+    try:
+        tokens = shlex.split(command, posix=True)
+    except ValueError:
+        match = _DB_CLI_WORD.search(command)
+        if match:
+            return [_Segment(cli=match.group(1).lower(), sqls=(), raw=True)]
+        return []
+
+    found: list[_Segment] = []
+    for raw_seg in _split_segments(tokens):
+        stripped = _skip_wrappers(raw_seg)
+        if not stripped:
+            continue
+        cli = _basename(stripped[0])
+        if cli not in DB_CLIS:
+            continue
+        sqls, incomplete = _extract_from_cli(cli, stripped)
+        if sqls and not incomplete:
+            found.append(_Segment(cli=cli, sqls=tuple(sqls), raw=False))
+        else:
+            found.append(_Segment(cli=cli, sqls=(), raw=True))
+    return found
+
+
+def _emit_refuse(
+    stream: TextIO,
+    *,
+    rule_id: str,
+    reason: str,
+    extra: str | None = None,
+) -> None:
+    stream.write("BLOCKED\n")
+    stream.write(f"Rule: {rule_id}\n")
+    stream.write(f"Reason: {reason}\n")
+    if extra:
+        stream.write(extra.rstrip() + "\n")
+
+
+def _emit_decision(stream: TextIO, decision: Decision) -> None:
+    extra = None
+    if decision.action == ACTION_APPROVAL:
+        extra = (
+            "REQUIRE_APPROVAL is refused by the hook so agents cannot silently write. "
+            "Use sql-write-gate check|exec with a human in the loop."
+        )
+    _emit_refuse(
+        stream,
+        rule_id=decision.rule_id,
+        reason=decision.reason,
+        extra=extra,
+    )
+
+
+def _raw_cli_reason(cli: str) -> str:
+    return (
+        f"Raw {cli} is blocked (interactive or unparseable; no -c/--command SQL). "
+        "Use sql-write-gate check|exec. Do not open a raw DB shell."
+    )
+
+
+def _gate(
+    *,
+    database: str | None,
+    db: str | None,
+    catalog: str | None,
+    policy: str | None,
+    agent: str,
+) -> WriteGate:
+    return WriteGate(
+        db_path=Path(db) if db else None,
+        database=database,
+        catalog_path=Path(catalog) if catalog else None,
+        policy_path=Path(policy) if policy else None,
+        agent=agent,
+    )
+
+
+def _check_sql(
+    sql: str,
+    *,
+    database: str | None,
+    db: str | None,
+    catalog: str | None,
+    policy: str | None,
+    agent: str,
+    err: TextIO,
+) -> int:
+    with _gate(
+        database=database,
+        db=db,
+        catalog=catalog,
+        policy=policy,
+        agent=agent,
+    ) as gate:
+        decision = gate.check(sql)
+    if decision.action == ACTION_ALLOW:
+        return 0
+    _emit_decision(err, decision)
+    return 2
+
+
+def run_hook(
+    *,
+    bash_command: str | None = None,
+    sql: str | None = None,
+    stdin_text: str | None = None,
+    database: str | None = None,
+    db: str | None = None,
+    catalog: str | None = None,
+    policy: str | None = None,
+    agent: str = "hook",
+    err: TextIO | None = None,
+) -> int:
+    """Evaluate a PreToolUse payload or test flags. Never executes user SQL.
+
+    Returns 0 (allow / not a DB CLI) or 2 (BLOCK / REQUIRE_APPROVAL / raw CLI).
+    """
+    err = err if err is not None else sys.stderr
+
+    if sql and sql.strip():
+        return _check_sql(
+            sql,
+            database=database,
+            db=db,
+            catalog=catalog,
+            policy=policy,
+            agent=agent,
+            err=err,
+        )
+
+    command = bash_command
+    if not command:
+        payload = parse_stdin_payload(stdin_text or "")
+        command = command_from_payload(payload) if payload is not None else None
+
+    if not command or not str(command).strip():
+        return 0
+
+    segments = inspect_bash(command)
+    if not segments:
+        return 0
+
+    for seg in segments:
+        if seg.raw or not seg.sqls:
+            _emit_refuse(
+                err,
+                rule_id=RULE_RAW_DB_CLI,
+                reason=_raw_cli_reason(seg.cli or "db-cli"),
+            )
+            return 2
+        for statement in seg.sqls:
+            rc = _check_sql(
+                statement,
+                database=database,
+                db=db,
+                catalog=catalog,
+                policy=policy,
+                agent=agent,
+                err=err,
+            )
+            if rc != 0:
+                return rc
+    return 0