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

write_gate/adapters/base.py

@@ -0,0 +1,101 @@
+"""Adapter routing: postgres:// / mysql:// / sqlite:// URLs vs DuckDB file paths."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Mapping
+
+from write_gate.paths import DB_PATH
+
+BACKEND_DUCKDB = "duckdb"
+BACKEND_POSTGRES = "postgres"
+BACKEND_MYSQL = "mysql"
+BACKEND_SQLITE = "sqlite"
+
+_PG_PREFIXES = ("postgres://", "postgresql://")
+_MYSQL_PREFIXES = ("mysql://", "mysql+pymysql://")
+_SQLITE_PREFIXES = ("sqlite:///", "sqlite+aiosqlite://")
+
+
+def is_postgres_url(value: str | None) -> bool:
+    if not value:
+        return False
+    lowered = value.strip().lower()
+    return lowered.startswith(_PG_PREFIXES)
+
+
+def is_mysql_url(value: str | None) -> bool:
+    if not value:
+        return False
+    lowered = value.strip().lower()
+    return lowered.startswith(_MYSQL_PREFIXES)
+
+
+def is_sqlite_url(value: str | None) -> bool:
+    if not value:
+        return False
+    lowered = value.strip().lower()
+    return lowered.startswith(_SQLITE_PREFIXES)
+
+
+def detect_backend(target: str | None) -> str:
+    if is_sqlite_url(target):
+        return BACKEND_SQLITE
+    if is_mysql_url(target):
+        return BACKEND_MYSQL
+    if is_postgres_url(target):
+        return BACKEND_POSTGRES
+    return BACKEND_DUCKDB
+
+
+def sqlglot_dialect(backend: str) -> str:
+    if backend in {BACKEND_POSTGRES, "postgresql", "pg"}:
+        return "postgres"
+    if backend == BACKEND_MYSQL:
+        return "mysql"
+    if backend == BACKEND_SQLITE:
+        return "sqlite"
+    return "duckdb"
+
+
+def count_sql(table: str, predicate: str | None, *, backend: str = BACKEND_DUCKDB) -> str:
+    """SELECT COUNT(*) of rows matching the write predicate.
+
+    Same COUNT form for DuckDB, Postgres, MySQL, and SQLite (EXPLAIN is optional elsewhere).
+    `backend` is accepted so callers can be explicit; quoting is ANSI identifiers.
+    """
+    del backend  # COUNT SQL is shared; dialect only affects WHERE rendering.
+    sql = f'SELECT COUNT(*) FROM "{table}"'
+    if predicate:
+        sql = f"{sql} WHERE {predicate}"
+    return sql
+
+
+def resolve_target(
+    *,
+    database: str | None = None,
+    database_url: str | None = None,
+    db_path: str | Path | None = None,
+    environ: Mapping[str, str] | None = None,
+) -> tuple[str, str]:
+    """Return (backend, target).
+
+    Priority:
+      1. ``database=``
+      2. ``database_url=``
+      3. ``db_path=`` (explicit path wins over env so tests/demo stay DuckDB)
+      4. ``DATABASE_URL`` env
+      5. default DuckDB warehouse
+    """
+    env = os.environ if environ is None else environ
+    for candidate in (database, database_url):
+        if candidate:
+            return detect_backend(candidate), candidate
+    if db_path is not None:
+        target = str(db_path)
+        return detect_backend(target), target
+    env_url = env.get("DATABASE_URL")
+    if env_url:
+        return detect_backend(env_url), env_url
+    return BACKEND_DUCKDB, str(DB_PATH)

write_gate/adapters/duckdb.py

@@ -0,0 +1,41 @@
+"""DuckDB warehouse helpers. User SQL must not be executed from here; use WriteGate."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import duckdb
+
+from write_gate.adapters.base import BACKEND_DUCKDB, count_sql as _count_sql
+from write_gate.paths import DB_PATH
+
+DIALECT = "duckdb"
+BACKEND = BACKEND_DUCKDB
+
+ORDERS_DDL = """
+CREATE TABLE orders (
+    order_id INTEGER PRIMARY KEY,
+    user_id INTEGER NOT NULL,
+    amount DOUBLE NOT NULL,
+    dt DATE NOT NULL,
+    email VARCHAR,
+    phone VARCHAR,
+    status VARCHAR NOT NULL
+)
+"""
+
+
+def count_sql(table: str, predicate: str | None) -> str:
+    """Blast-radius estimate: COUNT(*) of the matching predicate."""
+    return _count_sql(table, predicate, backend=BACKEND_DUCKDB)
+
+
+def connect(db_path: Path | None = None, *, read_only: bool = False) -> duckdb.DuckDBPyConnection:
+    path = Path(db_path) if db_path else DB_PATH
+    path.parent.mkdir(parents=True, exist_ok=True)
+    return duckdb.connect(str(path), read_only=read_only)
+
+
+def execute_user_sql(conn: duckdb.DuckDBPyConnection, sql: str):
+    """Run already-gated SQL. Called only by write_gate.wrapper.WriteGate."""
+    return conn.execute(sql)

write_gate/adapters/mysql.py

@@ -0,0 +1,115 @@
+"""MySQL adapter. User SQL must not be executed from here; use WriteGate."""
+
+from __future__ import annotations
+
+from typing import Any
+from urllib.parse import unquote, urlparse
+
+from write_gate.adapters.base import BACKEND_MYSQL, count_sql as _count_sql
+
+DIALECT = "mysql"
+BACKEND = BACKEND_MYSQL
+
+ORDERS_DDL = """
+CREATE TABLE orders (
+    order_id INTEGER PRIMARY KEY,
+    user_id INTEGER NOT NULL,
+    amount DOUBLE NOT NULL,
+    dt DATE NOT NULL,
+    email VARCHAR(255),
+    phone VARCHAR(64),
+    status VARCHAR(64) NOT NULL
+)
+"""
+
+
+def count_sql(table: str, predicate: str | None) -> str:
+    """Blast-radius estimate: COUNT(*) of the matching predicate (not EXPLAIN)."""
+    return _count_sql(table, predicate, backend=BACKEND_MYSQL)
+
+
+def explain_sql(table: str, predicate: str | None) -> str:
+    """Optional EXPLAIN form. Blast-radius uses count_sql() instead."""
+    return f"EXPLAIN {count_sql(table, predicate)}"
+
+
+def _parse_mysql_dsn(dsn: str) -> dict[str, Any]:
+    """Parse mysql:// or mysql+pymysql:// into driver connect kwargs."""
+    parsed = urlparse(dsn.strip())
+    database = unquote(parsed.path.lstrip("/")) if parsed.path else ""
+    kwargs: dict[str, Any] = {
+        "host": parsed.hostname or "localhost",
+        "port": parsed.port or 3306,
+        "user": unquote(parsed.username) if parsed.username else None,
+        "password": unquote(parsed.password) if parsed.password is not None else "",
+        "database": database or None,
+    }
+    return {k: v for k, v in kwargs.items() if v is not None}
+
+
+def _connect_raw(dsn: str, **kwargs: Any) -> Any:
+    """Open a driver connection. Tests may monkeypatch this."""
+    timeout = kwargs.pop("connect_timeout", 3)
+    params = _parse_mysql_dsn(dsn)
+    params.update(kwargs)
+    params.setdefault("connect_timeout", timeout)
+    try:
+        import pymysql
+
+        return pymysql.connect(**params)
+    except ImportError:
+        pass
+    try:
+        import mysql.connector
+
+        # mysql.connector uses connection_timeout, not connect_timeout.
+        if "connect_timeout" in params:
+            params["connection_timeout"] = params.pop("connect_timeout")
+        return mysql.connector.connect(**params)
+    except ImportError as exc:
+        raise ImportError(
+            "MySQL support requires pymysql. "
+            "Install with: pip install 'sql-write-gate[mysql]'"
+        ) from exc
+
+
+class MySQLConnection:
+    """DuckDB-like execute/fetchone surface over pymysql or mysql.connector."""
+
+    def __init__(self, raw: Any) -> None:
+        self._raw = raw
+        try:
+            self._raw.autocommit = True
+        except Exception:
+            pass
+
+    def execute(self, sql: str):
+        cursor_fn = getattr(self._raw, "cursor", None)
+        if callable(cursor_fn):
+            cur = cursor_fn()
+            cur.execute(sql)
+            return cur
+        execute = getattr(self._raw, "execute", None)
+        if callable(execute):
+            return execute(sql)
+        raise RuntimeError("MySQL connection does not support execute/cursor")
+
+    def close(self) -> None:
+        close = getattr(self._raw, "close", None)
+        if callable(close):
+            close()
+
+    def commit(self) -> None:
+        commit = getattr(self._raw, "commit", None)
+        if callable(commit):
+            commit()
+
+
+def connect(dsn: str, *, read_only: bool = False) -> MySQLConnection:
+    del read_only
+    return MySQLConnection(_connect_raw(dsn))
+
+
+def execute_user_sql(conn: Any, sql: str):
+    """Run already-gated SQL. Called only by write_gate.wrapper.WriteGate."""
+    return conn.execute(sql)

write_gate/adapters/postgres.py

@@ -0,0 +1,96 @@
+"""PostgreSQL adapter. User SQL must not be executed from here; use WriteGate."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from write_gate.adapters.base import BACKEND_POSTGRES, count_sql as _count_sql
+
+DIALECT = "postgres"
+BACKEND = BACKEND_POSTGRES
+
+ORDERS_DDL = """
+CREATE TABLE orders (
+    order_id INTEGER PRIMARY KEY,
+    user_id INTEGER NOT NULL,
+    amount DOUBLE PRECISION NOT NULL,
+    dt DATE NOT NULL,
+    email VARCHAR,
+    phone VARCHAR,
+    status VARCHAR NOT NULL
+)
+"""
+
+
+def count_sql(table: str, predicate: str | None) -> str:
+    """Blast-radius estimate: COUNT(*) of the matching predicate (not EXPLAIN)."""
+    return _count_sql(table, predicate, backend=BACKEND_POSTGRES)
+
+
+def explain_sql(table: str, predicate: str | None) -> str:
+    """Optional EXPLAIN form. Blast-radius uses count_sql() instead."""
+    return f"EXPLAIN {count_sql(table, predicate)}"
+
+
+def _connect_raw(dsn: str, **kwargs: Any) -> Any:
+    """Open a driver connection. Tests may monkeypatch this."""
+    timeout = kwargs.pop("connect_timeout", 3)
+    try:
+        import psycopg
+
+        return psycopg.connect(dsn, autocommit=True, connect_timeout=timeout, **kwargs)
+    except ImportError:
+        pass
+    try:
+        import psycopg2
+
+        conn = psycopg2.connect(dsn, connect_timeout=timeout, **kwargs)
+        conn.autocommit = True
+        return conn
+    except ImportError as exc:
+        raise ImportError(
+            "PostgreSQL support requires psycopg. "
+            "Install with: pip install 'sql-write-gate[postgres]'"
+        ) from exc
+
+
+class PostgresConnection:
+    """DuckDB-like execute/fetchone surface over psycopg or psycopg2."""
+
+    def __init__(self, raw: Any) -> None:
+        self._raw = raw
+        try:
+            self._raw.autocommit = True
+        except Exception:
+            pass
+
+    def execute(self, sql: str):
+        cursor_fn = getattr(self._raw, "cursor", None)
+        if callable(cursor_fn):
+            cur = cursor_fn()
+            cur.execute(sql)
+            return cur
+        execute = getattr(self._raw, "execute", None)
+        if callable(execute):
+            return execute(sql)
+        raise RuntimeError("PostgreSQL connection does not support execute/cursor")
+
+    def close(self) -> None:
+        close = getattr(self._raw, "close", None)
+        if callable(close):
+            close()
+
+    def commit(self) -> None:
+        commit = getattr(self._raw, "commit", None)
+        if callable(commit):
+            commit()
+
+
+def connect(dsn: str, *, read_only: bool = False) -> PostgresConnection:
+    del read_only
+    return PostgresConnection(_connect_raw(dsn))
+
+
+def execute_user_sql(conn: Any, sql: str):
+    """Run already-gated SQL. Called only by write_gate.wrapper.WriteGate."""
+    return conn.execute(sql)

write_gate/adapters/sqlite.py

@@ -0,0 +1,116 @@
+"""SQLite adapter. User SQL must not be executed from here; use WriteGate."""
+
+from __future__ import annotations
+
+import sqlite3
+from pathlib import Path
+from typing import Any
+from urllib.parse import unquote
+
+from write_gate.adapters.base import BACKEND_SQLITE, count_sql as _count_sql
+
+DIALECT = "sqlite"
+BACKEND = BACKEND_SQLITE
+
+ORDERS_DDL = """
+CREATE TABLE orders (
+    order_id INTEGER PRIMARY KEY,
+    user_id INTEGER NOT NULL,
+    amount REAL NOT NULL,
+    dt DATE NOT NULL,
+    email TEXT,
+    phone TEXT,
+    status TEXT NOT NULL
+)
+"""
+
+
+def count_sql(table: str, predicate: str | None) -> str:
+    """Blast-radius estimate: COUNT(*) of the matching predicate (not EXPLAIN)."""
+    return _count_sql(table, predicate, backend=BACKEND_SQLITE)
+
+
+def explain_sql(table: str, predicate: str | None) -> str:
+    """Optional EXPLAIN form. Blast-radius uses count_sql() instead."""
+    return f"EXPLAIN {count_sql(table, predicate)}"
+
+
+def _parse_sqlite_path(dsn: str) -> str:
+    """Extract a filesystem path (or :memory:) from sqlite:/// / sqlite+aiosqlite://."""
+    raw = dsn.strip()
+    lower = raw.lower()
+    if lower.startswith("sqlite+aiosqlite://"):
+        rest = raw[len("sqlite+aiosqlite://") :]
+    elif lower.startswith("sqlite://"):
+        rest = raw[len("sqlite://") :]
+    else:
+        raise ValueError(f"not a sqlite URL: {dsn}")
+
+    if "?" in rest:
+        rest = rest.split("?", 1)[0]
+    rest = unquote(rest)
+
+    if rest in (":memory:", "/:memory:"):
+        return ":memory:"
+    # Four-slash absolute: //tmp/x.db -> /tmp/x.db
+    if rest.startswith("//"):
+        return rest[1:]
+    # Three-slash form: /tmp/x.db or /relative.db
+    if rest.startswith("/"):
+        return rest
+    return rest or ":memory:"
+
+
+def _connect_raw(dsn: str, **kwargs: Any) -> Any:
+    """Open a stdlib sqlite3 connection. Tests may monkeypatch this."""
+    path = _parse_sqlite_path(dsn)
+    timeout = kwargs.pop("timeout", 3.0)
+    if path != ":memory:":
+        parent = Path(path).parent
+        if str(parent) not in ("", "."):
+            parent.mkdir(parents=True, exist_ok=True)
+    try:
+        conn = sqlite3.connect(path, timeout=timeout, **kwargs)
+    except TypeError:
+        # Unexpected kwargs from callers; retry with path only.
+        conn = sqlite3.connect(path, timeout=timeout)
+    conn.isolation_level = None  # autocommit-like
+    return conn
+
+
+class SQLiteConnection:
+    """DuckDB-like execute/fetchone surface over stdlib sqlite3."""
+
+    def __init__(self, raw: Any) -> None:
+        self._raw = raw
+
+    def execute(self, sql: str):
+        cursor_fn = getattr(self._raw, "cursor", None)
+        if callable(cursor_fn):
+            cur = cursor_fn()
+            cur.execute(sql)
+            return cur
+        execute = getattr(self._raw, "execute", None)
+        if callable(execute):
+            return execute(sql)
+        raise RuntimeError("SQLite connection does not support execute/cursor")
+
+    def close(self) -> None:
+        close = getattr(self._raw, "close", None)
+        if callable(close):
+            close()
+
+    def commit(self) -> None:
+        commit = getattr(self._raw, "commit", None)
+        if callable(commit):
+            commit()
+
+
+def connect(dsn: str, *, read_only: bool = False) -> SQLiteConnection:
+    del read_only
+    return SQLiteConnection(_connect_raw(dsn))
+
+
+def execute_user_sql(conn: Any, sql: str):
+    """Run already-gated SQL. Called only by write_gate.wrapper.WriteGate."""
+    return conn.execute(sql)

write_gate/approvals.py

@@ -0,0 +1,205 @@
+"""JSONL approval queue. REQUIRE_APPROVAL SQL is recorded and not executed."""
+
+from __future__ import annotations
+
+import json
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from write_gate.decision import Decision
+from write_gate.paths import APPROVALS_PATH, LOG_DIR
+
+STATUS_PENDING = "pending"
+STATUS_APPROVED = "approved"
+STATUS_REJECTED = "rejected"
+
+_ID_LEN = 12
+
+
+class ApprovalError(Exception):
+    """Missing, not pending, or invalid approval record."""
+
+
+@dataclass
+class ApprovalRecord:
+    id: str
+    status: str
+    sql: str
+    database: str | None = None
+    db_path: str | None = None
+    policy_path: str | None = None
+    catalog_path: str | None = None
+    created_at: str = ""
+    decision: dict[str, Any] = field(default_factory=dict)
+    backend: str | None = None
+    agent: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "status": self.status,
+            "sql": self.sql,
+            "database": self.database,
+            "db_path": self.db_path,
+            "policy_path": self.policy_path,
+            "catalog_path": self.catalog_path,
+            "created_at": self.created_at,
+            "decision": self.decision,
+            "backend": self.backend,
+            "agent": self.agent,
+        }
+
+    @classmethod
+    def from_dict(cls, raw: dict[str, Any]) -> "ApprovalRecord":
+        decision = raw.get("decision")
+        if not isinstance(decision, dict):
+            decision = {}
+        return cls(
+            id=str(raw.get("id") or ""),
+            status=str(raw.get("status") or STATUS_PENDING),
+            sql=str(raw.get("sql") or ""),
+            database=raw.get("database"),
+            db_path=raw.get("db_path"),
+            policy_path=raw.get("policy_path"),
+            catalog_path=raw.get("catalog_path"),
+            created_at=str(raw.get("created_at") or ""),
+            decision=decision,
+            backend=raw.get("backend"),
+            agent=raw.get("agent"),
+        )
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _new_id() -> str:
+    return uuid.uuid4().hex[:_ID_LEN]
+
+
+def _path(path: Path | str | None = None) -> Path:
+    return Path(path) if path else APPROVALS_PATH
+
+
+def _load(path: Path) -> dict[str, dict[str, Any]]:
+    if not path.exists():
+        return {}
+    records: dict[str, dict[str, Any]] = {}
+    try:
+        text = path.read_text(encoding="utf-8")
+    except OSError:
+        return {}
+    for line in text.splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            raw = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        if not isinstance(raw, dict):
+            continue
+        rec_id = str(raw.get("id") or "")
+        if rec_id:
+            records[rec_id] = raw
+    return records
+
+
+def _save(path: Path, records: dict[str, dict[str, Any]]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    body = "".join(json.dumps(rec, ensure_ascii=False) + "\n" for rec in records.values())
+    tmp.write_text(body, encoding="utf-8")
+    tmp.replace(path)
+
+
+def get_approval(approval_id: str, path: Path | str | None = None) -> ApprovalRecord | None:
+    dest = _path(path)
+    raw = _load(dest).get(str(approval_id))
+    if not raw:
+        return None
+    return ApprovalRecord.from_dict(raw)
+
+
+def list_pending(path: Path | str | None = None) -> list[ApprovalRecord]:
+    dest = _path(path)
+    pending: list[ApprovalRecord] = []
+    for raw in _load(dest).values():
+        rec = ApprovalRecord.from_dict(raw)
+        if rec.status == STATUS_PENDING and rec.id:
+            pending.append(rec)
+    pending.sort(key=lambda r: r.created_at)
+    return pending
+
+
+def enqueue_approval(
+    *,
+    sql: str,
+    decision: Decision,
+    database: str | None = None,
+    db_path: str | None = None,
+    policy_path: str | None = None,
+    catalog_path: str | None = None,
+    path: Path | str | None = None,
+    backend: str | None = None,
+    agent: str | None = None,
+) -> ApprovalRecord:
+    dest = _path(path)
+    records = _load(dest)
+    approval_id = _new_id()
+    while approval_id in records:
+        approval_id = _new_id()
+    rec = ApprovalRecord(
+        id=approval_id,
+        status=STATUS_PENDING,
+        sql=sql,
+        database=database,
+        db_path=db_path,
+        policy_path=str(policy_path) if policy_path else None,
+        catalog_path=str(catalog_path) if catalog_path else None,
+        created_at=_now(),
+        decision=decision.to_dict(),
+        backend=backend,
+        agent=agent,
+    )
+    # Snapshot includes approval_id once attached on the live decision.
+    rec.decision = {**rec.decision, "approval_id": approval_id}
+    records[approval_id] = rec.to_dict()
+    _save(dest, records)
+    return rec
+
+
+def set_status(
+    approval_id: str,
+    status: str,
+    path: Path | str | None = None,
+) -> ApprovalRecord:
+    if status not in {STATUS_PENDING, STATUS_APPROVED, STATUS_REJECTED}:
+        raise ApprovalError(f"invalid approval status: {status}")
+    dest = _path(path)
+    records = _load(dest)
+    raw = records.get(str(approval_id))
+    if not raw:
+        raise ApprovalError(f"approval not found: {approval_id}")
+    rec = ApprovalRecord.from_dict(raw)
+    if rec.status != STATUS_PENDING:
+        raise ApprovalError(f"approval not pending: {approval_id}")
+    rec.status = status
+    records[rec.id] = rec.to_dict()
+    _save(dest, records)
+    return rec
+
+
+def mark_approved(approval_id: str, path: Path | str | None = None) -> ApprovalRecord:
+    return set_status(approval_id, STATUS_APPROVED, path=path)
+
+
+def mark_rejected(approval_id: str, path: Path | str | None = None) -> ApprovalRecord:
+    return set_status(approval_id, STATUS_REJECTED, path=path)
+
+
+def ensure_log_dir() -> None:
+    LOG_DIR.mkdir(parents=True, exist_ok=True)