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

write_gate/init.py

@@ -0,0 +1,84 @@
+"""Scaffold a starter sql-write-gate project (policy, catalog, getting started)."""
+
+from __future__ import annotations
+
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+
+TEMPLATES_DIR = Path(__file__).resolve().parent / "templates"
+
+# (template filename, destination relative path)
+STARTER_FILES: tuple[tuple[str, str], ...] = (
+    ("policy.yaml", "policy.yaml"),
+    ("catalog.json", "catalog.json"),
+    ("GETTING_STARTED.md", "GETTING_STARTED.md"),
+)
+
+
+@dataclass(frozen=True)
+class InitResult:
+    created: list[str]
+    skipped: list[str]
+    overwritten: list[str]
+
+    @property
+    def all_paths(self) -> list[str]:
+        return [*self.created, *self.skipped, *self.overwritten]
+
+
+def templates_dir() -> Path:
+    """Resolve packaged templates (editable install or wheel)."""
+    path = TEMPLATES_DIR
+    if not path.is_dir():
+        raise FileNotFoundError(f"sql-write-gate templates missing: {path}")
+    return path
+
+
+def init_project(target_dir: Path | str = ".", *, force: bool = False) -> InitResult:
+    """Create starter files under *target_dir*.
+
+    Without *force*, existing files are left unchanged and listed as skipped.
+    With *force*, existing files are overwritten.
+    """
+    dest_root = Path(target_dir).expanduser().resolve()
+    dest_root.mkdir(parents=True, exist_ok=True)
+    src_root = templates_dir()
+
+    created: list[str] = []
+    skipped: list[str] = []
+    overwritten: list[str] = []
+
+    for template_name, rel in STARTER_FILES:
+        src = src_root / template_name
+        if not src.is_file():
+            raise FileNotFoundError(f"template not found: {src}")
+        dest = dest_root / rel
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        label = str(dest.relative_to(dest_root)) if dest.is_relative_to(dest_root) else str(dest)
+
+        if dest.exists() and not force:
+            skipped.append(label)
+            continue
+        if dest.exists() and force:
+            shutil.copyfile(src, dest)
+            overwritten.append(label)
+        else:
+            shutil.copyfile(src, dest)
+            created.append(label)
+
+    return InitResult(created=created, skipped=skipped, overwritten=overwritten)
+
+
+def format_init_report(result: InitResult, target_dir: Path | str) -> str:
+    root = Path(target_dir).expanduser().resolve()
+    lines = [f"sql-write-gate init → {root}"]
+    for path in result.created:
+        lines.append(f"  created: {path}")
+    for path in result.overwritten:
+        lines.append(f"  overwritten: {path}")
+    for path in result.skipped:
+        lines.append(f"  skipped (exists): {path}")
+    if not result.created and not result.overwritten and result.skipped:
+        lines.append("  (nothing written; use --force to overwrite)")
+    return "\n".join(lines) + "\n"

write_gate/mcp_server.py

@@ -0,0 +1,71 @@
+"""MCP stdio server. FastMCP is lazy-imported so mcp_tools tests need no SDK."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def _fastmcp():
+    try:
+        from mcp.server.fastmcp import FastMCP
+    except ImportError as exc:  # optional extra
+        raise ImportError('pip install -e ".[mcp]"') from exc
+    return FastMCP
+
+
+def create_server(
+    *,
+    database: str | None = None,
+    db: str | None = None,
+    catalog: str | None = None,
+    policy: str | None = None,
+    agent: str = "mcp",
+):
+    """Build a FastMCP server exposing query_sql and write_sql (ALLOW executes)."""
+    from write_gate.mcp_tools import query_sql as check_query
+    from write_gate.mcp_tools import write_sql as check_write
+
+    FastMCP = _fastmcp()
+    mcp = FastMCP("sql-write-gate")
+    gate_kwargs = {
+        "database": database,
+        "db_path": db,
+        "catalog_path": catalog,
+        "policy_path": policy,
+        "agent": agent or "mcp",
+    }
+
+    @mcp.tool()
+    def query_sql(sql: str) -> dict[str, Any]:
+        """Evaluate a SELECT (or any SQL) through sql-write-gate. ALLOW executes."""
+        return check_query(sql, **gate_kwargs)
+
+    @mcp.tool()
+    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)
+
+    return mcp
+
+
+def run_server(
+    *,
+    database: str | None = None,
+    db: str | None = None,
+    catalog: str | None = None,
+    policy: str | None = None,
+    agent: str = "mcp",
+) -> None:
+    """Start the MCP server on stdio (mcp.run())."""
+    mcp = create_server(
+        database=database,
+        db=db,
+        catalog=catalog,
+        policy=policy,
+        agent=agent,
+    )
+    mcp.run()
+
+
+if __name__ == "__main__":
+    run_server()

write_gate/mcp_tools.py

@@ -0,0 +1,197 @@
+"""MCP tool functions: execute ALLOW SQL via WriteGate.
+
+query_sql / write_sql call WriteGate.execute (not raw DuckDB). BLOCK and
+REQUIRE_APPROVAL already return None and do not write. Honor database= /
+DATABASE_URL the same way as the CLI.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from write_gate.decision import ACTION_ALLOW
+from write_gate.wrapper import WriteGate
+
+QUERY_ROW_CAP = 50
+
+
+def _as_path(value: str | Path | None) -> Path | None:
+    if value is None:
+        return None
+    return Path(value)
+
+
+def _gate(
+    *,
+    database: str | None = None,
+    db_path: str | Path | None = None,
+    catalog_path: str | Path | None = None,
+    policy_path: str | Path | None = None,
+    approvals_path: str | Path | None = None,
+    agent: str = "mcp",
+) -> WriteGate:
+    return WriteGate(
+        database=database,
+        db_path=_as_path(db_path),
+        catalog_path=_as_path(catalog_path),
+        policy_path=_as_path(policy_path),
+        approvals_path=_as_path(approvals_path),
+        agent=agent or "mcp",
+    )
+
+
+def _json_cell(value: Any) -> Any:
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    if isinstance(value, bytes):
+        return value.decode("utf-8", errors="replace")
+    iso = getattr(value, "isoformat", None)
+    if callable(iso):
+        return iso()
+    return str(value)
+
+
+def _fetch_rows(result: Any, *, cap: int) -> list[list[Any]]:
+    if result is None:
+        return []
+    fetched = None
+    fetchmany = getattr(result, "fetchmany", None)
+    if callable(fetchmany):
+        try:
+            fetched = fetchmany(cap)
+        except Exception:
+            fetched = None
+    if fetched is None:
+        fetchall = getattr(result, "fetchall", None)
+        if not callable(fetchall):
+            return []
+        try:
+            fetched = fetchall()[:cap]
+        except Exception:
+            return []
+    rows: list[list[Any]] = []
+    for row in fetched:
+        rows.append([_json_cell(c) for c in row])
+    return rows
+
+
+def _rowcount(result: Any) -> int | None:
+    if result is None:
+        return None
+    rc = getattr(result, "rowcount", None)
+    if isinstance(rc, int) and rc >= 0:
+        return rc
+    return None
+
+
+def decision_payload(
+    decision: Any,
+    *,
+    executed: bool = False,
+    rowcount: int | None = None,
+    rows: list[list[Any]] | None = None,
+) -> dict[str, Any]:
+    """JSON-serializable gate result for MCP tools / python -c demos."""
+    payload: dict[str, Any] = {
+        "action": decision.action,
+        "rule_id": decision.rule_id,
+        "reason": decision.reason,
+        "operation": decision.operation,
+        "table": decision.table,
+        "risk": decision.risk,
+        "executed": bool(executed),
+    }
+    if getattr(decision, "approval_id", None):
+        payload["approval_id"] = decision.approval_id
+    if rowcount is not None:
+        payload["rowcount"] = rowcount
+    if rows is not None:
+        payload["rows"] = rows
+    return payload
+
+
+def _execute(
+    sql: str,
+    *,
+    include_rows: bool,
+    database: str | None = None,
+    db_path: str | Path | None = None,
+    catalog_path: str | Path | None = None,
+    policy_path: str | Path | None = None,
+    approvals_path: str | Path | None = None,
+    agent: str = "mcp",
+) -> dict[str, Any]:
+    with _gate(
+        database=database,
+        db_path=db_path,
+        catalog_path=catalog_path,
+        policy_path=policy_path,
+        approvals_path=approvals_path,
+        agent=agent,
+    ) as gate:
+        decision, result = gate.execute(sql)
+        executed = decision.action == ACTION_ALLOW and result is not None
+        rows = None
+        rowcount = None
+        if executed:
+            rowcount = _rowcount(result)
+            if include_rows:
+                rows = _fetch_rows(result, cap=QUERY_ROW_CAP)
+                if rowcount is None:
+                    rowcount = len(rows)
+    return decision_payload(decision, executed=executed, rowcount=rowcount, rows=rows)
+
+
+def query_sql(
+    sql: str,
+    *,
+    database: str | None = None,
+    db_path: str | Path | None = None,
+    catalog_path: str | Path | None = None,
+    policy_path: str | Path | None = None,
+    approvals_path: str | Path | None = None,
+    agent: str = "mcp",
+) -> dict[str, Any]:
+    """Gate then (on ALLOW) execute SQL via WriteGate.execute.
+
+    SELECT results include a capped ``rows`` list. BLOCK / REQUIRE_APPROVAL
+    do not run user SQL. REQUIRE_APPROVAL is queued (approval_id).
+    """
+    return _execute(
+        sql,
+        include_rows=True,
+        database=database,
+        db_path=db_path,
+        catalog_path=catalog_path,
+        policy_path=policy_path,
+        approvals_path=approvals_path,
+        agent=agent,
+    )
+
+
+def write_sql(
+    sql: str,
+    *,
+    database: str | None = None,
+    db_path: str | Path | None = None,
+    catalog_path: str | Path | None = None,
+    policy_path: str | Path | None = None,
+    approvals_path: str | Path | None = None,
+    agent: str = "mcp",
+) -> dict[str, Any]:
+    """Gate then (on ALLOW) execute INSERT/UPDATE/DELETE/DDL via WriteGate.execute.
+
+    BLOCK / REQUIRE_APPROVAL return executed=false and do not write.
+    REQUIRE_APPROVAL is queued (approval_id); approve <id> later to write.
+    """
+    return _execute(
+        sql,
+        include_rows=False,
+        database=database,
+        db_path=db_path,
+        catalog_path=catalog_path,
+        policy_path=policy_path,
+        approvals_path=approvals_path,
+        agent=agent,
+    )