sqlguard-ninja 0.1.0__py3-none-any.whl

sqlguard_ninja/integrations/checkpoint.py

@@ -0,0 +1,342 @@
+"""Pre-execution SQL checkpoint for agents.
+
+The checkpoint separates two questions:
+
+1. *Policy evaluation* — what does the policy say about this SQL for this caller?
+   (:meth:`Guard.validate`, a pure function of SQL + policy + trusted context.)
+2. *Execution authorization* — given that verdict, the enforcement mode and any human
+   approval, may this exact SQL run now?
+
+Nothing an LLM produces can change either answer. The role and tenant come from a
+:class:`RequestContext` the application builds from its own authentication. Tool arguments,
+SQL comments and model "explanations" are not inputs to either decision. Approvals are bound
+to the SHA-256 of the exact SQL, the trace and the requesting principal, and can be used once.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+import uuid
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import Any, Protocol, TypeVar
+
+from sqlguard_ninja.core.findings import Decision, ValidationResult
+from sqlguard_ninja.core.guard import Guard
+from sqlguard_ninja.integrations.audit import AuditEvent, AuditSink, redact_sql
+
+T = TypeVar("T")
+
+
+class EnforcementMode(StrEnum):
+    BLOCK = "block"
+    """DENY and REVIEW are not executed."""
+    REVIEW = "review"
+    """DENY is not executed; REVIEW runs only with a valid human approval."""
+    AUDIT = "audit"
+    """Everything runs; findings are recorded with ``enforced=False``. Monitoring only."""
+
+
+@dataclass(frozen=True)
+class RequestContext:
+    """Trusted, server-side facts about the request. Build it from your auth layer only."""
+
+    principal: str
+    role: str | None
+    tenant_id: Any = None
+    trace_id: str = field(default_factory=lambda: uuid.uuid4().hex)
+    attributes: Mapping[str, Any] = field(default_factory=dict)
+
+    def policy_context(self, tenant_key: str) -> dict[str, Any]:
+        ctx = dict(self.attributes)
+        if self.tenant_id is not None:
+            ctx[tenant_key] = self.tenant_id
+        return ctx
+
+
+@dataclass(frozen=True)
+class ApprovalRequest:
+    request_id: str
+    trace_id: str
+    principal: str
+    role: str | None
+    sql_sha256: str
+    summary: str
+    created_at: float
+
+
+@dataclass(frozen=True)
+class ApprovalDecision:
+    approved: bool
+    approver: str
+    reason: str = ""
+
+
+class Approver(Protocol):
+    """Decide (or defer) an approval request. Return None while pending."""
+
+    def decide(self, request: ApprovalRequest) -> ApprovalDecision | None: ...
+
+
+class DenyAllApprover:
+    """Default: nothing is ever approved."""
+
+    def decide(self, request: ApprovalRequest) -> ApprovalDecision | None:
+        return ApprovalDecision(False, "system", "no approver configured")
+
+
+class CallbackApprover:
+    """Synchronous approval through a callback (e.g. a CLI prompt or chat-ops bot)."""
+
+    def __init__(self, fn: Callable[[ApprovalRequest], ApprovalDecision | None]) -> None:
+        self.fn = fn
+
+    def decide(self, request: ApprovalRequest) -> ApprovalDecision | None:
+        return self.fn(request)
+
+
+class ApprovalQueue:
+    """Asynchronous human approval.
+
+    A REVIEW query is parked as a pending request. A human (a different principal) approves it
+    out of band; the agent may then re-submit *the identical SQL* within ``ttl_s`` seconds.
+    Each approval is single-use.
+    """
+
+    def __init__(self, ttl_s: float = 900.0, clock: Callable[[], float] = time.time) -> None:
+        self.ttl_s = ttl_s
+        self.clock = clock
+        self.pending: dict[str, ApprovalRequest] = {}
+        self._decided: dict[str, tuple[ApprovalRequest, ApprovalDecision, float]] = {}
+        self._lock = threading.Lock()
+
+    @staticmethod
+    def _key(req: ApprovalRequest) -> str:
+        return f"{req.trace_id}|{req.principal}|{req.sql_sha256}"
+
+    def decide(self, request: ApprovalRequest) -> ApprovalDecision | None:
+        key = self._key(request)
+        with self._lock:
+            decided = self._decided.get(key)
+            if decided is not None:
+                _req, decision, at = decided
+                if self.clock() - at > self.ttl_s:
+                    del self._decided[key]
+                    return ApprovalDecision(False, decision.approver, "approval expired")
+                del self._decided[key]  # single use
+                return decision
+            if not any(self._key(p) == key for p in self.pending.values()):
+                self.pending[request.request_id] = request
+            return None
+
+    def _resolve(self, request_id: str, approver: str, approved: bool, reason: str) -> None:
+        with self._lock:
+            req = self.pending.get(request_id)
+            if req is None:
+                raise KeyError(f"no pending approval {request_id!r}")
+            if approver == req.principal:
+                raise PermissionError("the requesting principal cannot approve its own query")
+            del self.pending[request_id]
+            self._decided[self._key(req)] = (
+                req,
+                ApprovalDecision(approved, approver, reason),
+                self.clock(),
+            )
+
+    def approve(self, request_id: str, approver: str, reason: str = "") -> None:
+        self._resolve(request_id, approver, True, reason)
+
+    def reject(self, request_id: str, approver: str, reason: str = "") -> None:
+        self._resolve(request_id, approver, False, reason)
+
+
+@dataclass(frozen=True)
+class CheckpointOutcome:
+    authorized: bool
+    """May the exact SQL be executed now?"""
+    enforced: bool
+    """False in audit mode: a negative verdict was recorded but did not stop execution."""
+    would_block: bool
+    """True whenever the policy verdict was not ALLOW, regardless of mode."""
+    mode: EnforcementMode
+    result: ValidationResult
+    agent_message: str
+    """Text safe to return to the LLM. Never contains approval tokens or internal detail."""
+    approval_request: ApprovalRequest | None = None
+    approval: ApprovalDecision | None = None
+    attempt: int = 1
+    locked_out: bool = False
+    audit_event: AuditEvent | None = None
+
+
+class SQLBlockedError(PermissionError):
+    def __init__(self, outcome: CheckpointOutcome) -> None:
+        super().__init__(outcome.agent_message)
+        self.outcome = outcome
+
+
+class SQLCheckpoint:
+    """Evaluate and authorize SQL before execution. Thread-safe."""
+
+    def __init__(
+        self,
+        guard: Guard,
+        mode: EnforcementMode | str = EnforcementMode.BLOCK,
+        *,
+        approver: Approver | None = None,
+        audit_sink: AuditSink | None = None,
+        max_denied_attempts: int | None = 5,
+        dialect: str | None = None,
+        log_redacted_sql: bool = True,
+    ) -> None:
+        self.guard = guard
+        self.mode = EnforcementMode(mode)
+        self.approver: Approver = approver or DenyAllApprover()
+        self.audit_sink = audit_sink
+        self.max_denied_attempts = max_denied_attempts
+        self.dialect = dialect or guard.policy.dialect
+        self.log_redacted_sql = log_redacted_sql
+        self._attempts: dict[str, int] = {}
+        self._denials: dict[str, int] = {}
+        self._lock = threading.Lock()
+
+    # ------------------------------------------------------------------ evaluate
+    def evaluate(self, sql: str, context: RequestContext) -> CheckpointOutcome:
+        started = time.perf_counter()
+        policy_ctx = context.policy_context(self.guard.policy.tenant.context_key)
+        result = self.guard.validate(
+            sql, role=context.role, dialect=self.dialect, context=policy_ctx
+        )
+        with self._lock:
+            attempt = self._attempts.get(context.trace_id, 0) + 1
+            self._attempts[context.trace_id] = attempt
+            denials = self._denials.get(context.trace_id, 0)
+        locked = self.max_denied_attempts is not None and denials >= self.max_denied_attempts
+
+        would_block = result.decision is not Decision.ALLOW
+        approval_request: ApprovalRequest | None = None
+        approval: ApprovalDecision | None = None
+        if locked and self.mode is not EnforcementMode.AUDIT:
+            authorized, enforced = False, True
+        elif result.decision is Decision.ALLOW:
+            authorized, enforced = True, True
+        elif self.mode is EnforcementMode.AUDIT:
+            authorized, enforced = True, False
+        elif result.decision is Decision.REVIEW and self.mode is EnforcementMode.REVIEW:
+            approval_request = ApprovalRequest(
+                request_id=uuid.uuid4().hex,
+                trace_id=context.trace_id,
+                principal=context.principal,
+                role=context.role,
+                sql_sha256=result.sql_sha256,
+                summary="; ".join(f.message for f in result.blocking_findings)[:500],
+                created_at=time.time(),
+            )
+            approval = self.approver.decide(approval_request)
+            if approval is not None and approval.approver == context.principal:
+                approval = ApprovalDecision(False, approval.approver, "self-approval rejected")
+            authorized = approval is not None and approval.approved
+            enforced = True
+        else:
+            authorized, enforced = False, True
+
+        if not authorized and enforced:
+            with self._lock:
+                self._denials[context.trace_id] = self._denials.get(context.trace_id, 0) + 1
+
+        message = self._agent_message(
+            result, authorized, enforced, locked, approval_request, approval
+        )
+        event = AuditEvent(
+            trace_id=context.trace_id,
+            principal=context.principal,
+            role=result.role,
+            tenant=None if context.tenant_id is None else str(context.tenant_id),
+            mode=self.mode.value,
+            decision=result.decision.value,
+            authorized=authorized,
+            enforced=enforced,
+            would_block=would_block,
+            sql_sha256=result.sql_sha256,
+            sql_redacted=redact_sql(sql, self.dialect) if self.log_redacted_sql else None,
+            dialect=self.dialect,
+            policy=result.policy_name,
+            findings=[
+                {"rule_id": f.rule_id, "action": f.action.value, "category": f.category.value}
+                for f in result.findings
+            ],
+            approval=None
+            if approval is None
+            else {"approved": approval.approved, "approver": approval.approver},
+            attempt=attempt,
+            elapsed_ms=(time.perf_counter() - started) * 1000,
+        )
+        if self.audit_sink is not None:
+            self.audit_sink.emit(event)
+        return CheckpointOutcome(
+            authorized=authorized,
+            enforced=enforced,
+            would_block=would_block,
+            mode=self.mode,
+            result=result,
+            agent_message=message,
+            approval_request=approval_request,
+            approval=approval,
+            attempt=attempt,
+            locked_out=locked,
+            audit_event=event,
+        )
+
+    @staticmethod
+    def _agent_message(
+        result: ValidationResult,
+        authorized: bool,
+        enforced: bool,
+        locked: bool,
+        approval_request: ApprovalRequest | None,
+        approval: ApprovalDecision | None,
+    ) -> str:
+        if locked and not authorized:
+            return (
+                "Query blocked: too many policy violations in this session. "
+                "No further queries will run in this conversation."
+            )
+        reasons = "; ".join(
+            f.message + (f" Suggestion: {f.remediation}" if f.remediation else "")
+            for f in result.blocking_findings
+        )
+        if authorized and not enforced:
+            return "Query executed (audit mode). Policy findings were recorded."
+        if authorized:
+            return "Query permitted."
+        if approval_request is not None:
+            if approval is None:
+                return (
+                    "Query requires human approval and was not executed. "
+                    f"Reason: {reasons} Approval is handled outside this conversation."
+                )
+            return f"Query was not approved and was not executed. Reason: {reasons}"
+        return (
+            f"Query blocked by policy and not executed. Reason: {reasons} "
+            "This decision is enforced by the server and cannot be overridden by the assistant."
+        )
+
+    # ------------------------------------------------------------------- execute
+    def execute(
+        self,
+        sql: str,
+        context: RequestContext,
+        executor: Callable[[str, Mapping[str, Any]], T],
+    ) -> T:
+        """Evaluate, then run ``executor(sql, params)`` only if authorized.
+
+        ``params`` contains the trusted context values (e.g. ``tenant_id``) for bind
+        parameters. The SQL is executed exactly as validated.
+        """
+        outcome = self.evaluate(sql, context)
+        if not outcome.authorized:
+            raise SQLBlockedError(outcome)
+        params = context.policy_context(self.guard.policy.tenant.context_key)
+        return executor(sql, params)

sqlguard_ninja/integrations/executors.py

@@ -0,0 +1,64 @@
+"""Small, safe executors used by the integrations and examples.
+
+These are the *last* line of defence after the checkpoint: they open read-only sessions so
+that the database itself refuses writes even if a query were authorized by mistake (or
+deliberately, in audit mode). Use a least-privilege database account in production.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+import threading
+from collections.abc import Callable, Mapping
+from typing import Any
+
+Rows = list[dict[str, Any]]
+
+
+def _bound_params(sql: str, params: Mapping[str, Any]) -> dict[str, Any]:
+    # Only pass parameters the statement mentions; sqlite3 ignores extras for named style,
+    # but filtering keeps driver errors clear.
+    return {k: v for k, v in params.items() if f":{k}" in sql or f"%({k})s" in sql}
+
+
+def sqlite_executor(conn: sqlite3.Connection, max_rows: int = 200) -> Callable[..., Rows]:
+    """Execute on SQLite with ``PRAGMA query_only = ON`` (writes fail at the engine).
+
+    Calls are serialised with a lock so one connection can be shared by threaded servers;
+    open it with ``check_same_thread=False`` for that.
+    """
+    conn.execute("PRAGMA query_only = ON")
+    lock = threading.Lock()
+
+    def run(sql: str, params: Mapping[str, Any]) -> Rows:
+        with lock:
+            cur = conn.execute(sql, _bound_params(sql, params))
+            cols = [d[0] for d in cur.description or ()]
+            return [dict(zip(cols, row, strict=False)) for row in cur.fetchmany(max_rows)]
+
+    return run
+
+
+def psycopg_executor(
+    conninfo: str, max_rows: int = 200, statement_timeout_ms: int = 5000
+) -> Callable[..., Rows]:  # pragma: no cover - exercised by PostgreSQL integration tests
+    """Execute on PostgreSQL in a READ ONLY transaction with a statement timeout.
+
+    Use ``%(tenant_id)s`` style bind parameters in SQL sent to PostgreSQL.
+    """
+    import psycopg  # noqa: PLC0415 - optional dependency
+    from psycopg.rows import dict_row  # noqa: PLC0415
+
+    def run(sql: str, params: Mapping[str, Any]) -> Rows:
+        with (
+            psycopg.connect(conninfo, row_factory=dict_row) as conn,
+            conn.transaction(force_rollback=True),
+        ):
+            conn.execute("SET TRANSACTION READ ONLY")
+            conn.execute(
+                "SELECT set_config('statement_timeout', %s, true)", (str(statement_timeout_ms),)
+            )
+            cur = conn.execute(sql, _bound_params(sql, params))
+            return list(cur.fetchmany(max_rows))
+
+    return run

sqlguard_ninja/integrations/fastapi.py

@@ -0,0 +1,89 @@
+"""FastAPI integration: a router exposing validation and guarded execution.
+
+Authentication is the application's job. Pass a dependency that returns a
+:class:`RequestContext` built from your own auth (session, JWT, mTLS...). The request body
+can never set the role or tenant. Install with ``pip install "sqlguard-ninja[fastapi]"``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from typing import Any
+
+from fastapi import APIRouter, Depends, HTTPException
+from pydantic import BaseModel, ConfigDict, Field
+
+from sqlguard_ninja.integrations.checkpoint import RequestContext, SQLBlockedError, SQLCheckpoint
+
+
+class SQLRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    sql: str = Field(max_length=100_000)
+
+
+class FindingOut(BaseModel):
+    rule_id: str
+    action: str
+    message: str
+    remediation: str | None = None
+
+
+class ValidateResponse(BaseModel):
+    decision: str
+    authorized: bool
+    message: str
+    findings: list[FindingOut]
+    trace_id: str
+
+
+def create_router(
+    checkpoint: SQLCheckpoint,
+    context_dependency: Callable[..., RequestContext],
+    executor: Callable[[str, Mapping[str, Any]], Any] | None = None,
+    prefix: str = "/v1/sql",
+) -> APIRouter:
+    router = APIRouter(prefix=prefix, tags=["sql"])
+
+    @router.post("/validate", response_model=ValidateResponse)
+    def validate(
+        body: SQLRequest,
+        ctx: RequestContext = Depends(context_dependency),  # noqa: B008
+    ) -> ValidateResponse:
+        outcome = checkpoint.evaluate(body.sql, ctx)
+        return ValidateResponse(
+            decision=outcome.result.decision.value,
+            authorized=outcome.authorized,
+            message=outcome.agent_message,
+            findings=[
+                FindingOut(
+                    rule_id=f.rule_id,
+                    action=f.action.value,
+                    message=f.message,
+                    remediation=f.remediation,
+                )
+                for f in outcome.result.findings
+            ],
+            trace_id=ctx.trace_id,
+        )
+
+    if executor is not None:
+        run = executor
+
+        @router.post("/query")
+        def query(
+            body: SQLRequest,
+            ctx: RequestContext = Depends(context_dependency),  # noqa: B008
+        ) -> dict[str, Any]:
+            try:
+                rows = checkpoint.execute(body.sql, ctx, run)
+            except SQLBlockedError as exc:
+                raise HTTPException(
+                    status_code=403,
+                    detail={
+                        "decision": exc.outcome.result.decision.value,
+                        "message": exc.outcome.agent_message,
+                    },
+                ) from exc
+            return {"rows": rows, "trace_id": ctx.trace_id}
+
+    return router

sqlguard_ninja/integrations/langchain.py

@@ -0,0 +1,67 @@
+"""LangChain integration: a SQL tool whose execution is gated by :class:`SQLCheckpoint`.
+
+The tool's argument schema contains **only** the query. Role, tenant and approvals cannot be
+supplied by the model; they come from ``context_provider``, which the application binds to its
+authenticated request. Install with ``pip install "sqlguard-ninja[langchain]"``.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Callable, Mapping
+from typing import Any
+
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+
+from sqlguard_ninja.integrations.checkpoint import RequestContext, SQLBlockedError, SQLCheckpoint
+
+Executor = Callable[[str, Mapping[str, Any]], Any]
+
+
+class SQLQueryInput(BaseModel):
+    """Arguments the model may send. Extra fields (e.g. ``role``) are rejected."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    query: str = Field(description="A single read-only SQL SELECT statement.")
+
+
+class GuardedSQLTool(BaseTool):
+    """Run SQL only after the checkpoint authorizes it; return rows or the policy reason."""
+
+    name: str = "sql_db_query"
+    description: str = (
+        "Run a read-only SQL query against the database and return rows as JSON. "
+        "Queries are checked against a server-side data access policy before execution; "
+        "if a query is blocked, the response explains why. Use bind parameter :tenant_id "
+        "for the current tenant."
+    )
+    args_schema: type[BaseModel] = SQLQueryInput
+    raise_on_block: bool = False
+
+    _checkpoint: SQLCheckpoint = PrivateAttr()
+    _executor: Executor = PrivateAttr()
+    _context_provider: Callable[[], RequestContext] = PrivateAttr()
+
+    def __init__(
+        self,
+        checkpoint: SQLCheckpoint,
+        executor: Executor,
+        context_provider: Callable[[], RequestContext],
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._checkpoint = checkpoint
+        self._executor = executor
+        self._context_provider = context_provider
+
+    def _run(self, query: str, **_: Any) -> str:
+        context = self._context_provider()
+        try:
+            rows = self._checkpoint.execute(query, context, self._executor)
+        except SQLBlockedError as exc:
+            if self.raise_on_block:
+                raise
+            return exc.outcome.agent_message
+        return json.dumps(rows, default=str)

sqlguard_ninja/integrations/langgraph.py

@@ -0,0 +1,87 @@
+"""LangGraph integration: guard and execute nodes plus a router.
+
+Trust model:
+
+* The request context is read from ``config["configurable"][CONTEXT_KEY]``, which the
+  *application* sets when it invokes the graph. Graph nodes (including LLM nodes) can write
+  state, but they cannot change the run configuration.
+* :func:`make_execute_node` re-evaluates the SQL through the checkpoint. It never trusts an
+  ``authorized`` flag found in state, so a node that tampers with state gains nothing.
+
+Install with ``pip install "sqlguard-ninja[langgraph]"``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from typing import Any
+
+from langchain_core.runnables import RunnableConfig
+
+from sqlguard_ninja.integrations.checkpoint import RequestContext, SQLBlockedError, SQLCheckpoint
+
+CONTEXT_KEY = "sqlninja_context"
+
+State = Mapping[str, Any]
+Node = Callable[[State, RunnableConfig], dict[str, Any]]
+
+
+def _context(config: Any) -> RequestContext:
+    configurable = (config or {}).get("configurable", {}) if isinstance(config, Mapping) else {}
+    ctx = configurable.get(CONTEXT_KEY)
+    if not isinstance(ctx, RequestContext):
+        raise PermissionError(
+            f"missing trusted RequestContext in config['configurable'][{CONTEXT_KEY!r}]"
+        )
+    return ctx
+
+
+def make_guard_node(checkpoint: SQLCheckpoint, sql_key: str = "sql") -> Node:
+    """Node that evaluates ``state[sql_key]`` and records the verdict under ``guard``."""
+
+    # LangGraph injects the run config only when the parameter is annotated RunnableConfig.
+    def guard_node(state: State, config: RunnableConfig) -> dict[str, Any]:
+        outcome = checkpoint.evaluate(str(state.get(sql_key, "")), _context(config))
+        return {
+            "guard": {
+                "authorized": outcome.authorized,
+                "decision": outcome.result.decision.value,
+                "message": outcome.agent_message,
+                "rules": [f.rule_id for f in outcome.result.blocking_findings],
+                "locked_out": outcome.locked_out,
+            },
+            "attempts": int(state.get("attempts", 0)) + 1,
+        }
+
+    return guard_node
+
+
+def make_execute_node(
+    checkpoint: SQLCheckpoint,
+    executor: Callable[[str, Mapping[str, Any]], Any],
+    sql_key: str = "sql",
+) -> Node:
+    """Node that executes the SQL *after re-checking it*; writes ``rows`` or ``error``."""
+
+    def execute_node(state: State, config: RunnableConfig) -> dict[str, Any]:
+        try:
+            rows = checkpoint.execute(str(state.get(sql_key, "")), _context(config), executor)
+        except SQLBlockedError as exc:
+            return {"rows": None, "error": exc.outcome.agent_message}
+        return {"rows": rows, "error": None}
+
+    return execute_node
+
+
+def route_after_guard(max_attempts: int = 3) -> Callable[[State], str]:
+    """Router: ``"execute"`` if authorized, ``"revise"`` to retry, ``"end"`` when exhausted."""
+
+    def route(state: State) -> str:
+        guard = state.get("guard") or {}
+        if guard.get("authorized"):
+            return "execute"
+        if guard.get("locked_out") or int(state.get("attempts", 0)) >= max_attempts:
+            return "end"
+        return "revise"
+
+    return route

sqlguard_ninja/policy/__init__.py

@@ -0,0 +1,39 @@
+"""Policy schema and loading."""
+
+from sqlguard_ninja.policy.model import (
+    Access,
+    Catalog,
+    ColumnDef,
+    CostThresholds,
+    Defaults,
+    Disclosure,
+    FunctionPolicy,
+    LabelDef,
+    Limits,
+    Policy,
+    PolicyError,
+    Role,
+    Sensitivity,
+    TableDef,
+    TableGrant,
+    TenantConfig,
+)
+
+__all__ = [
+    "Access",
+    "Catalog",
+    "ColumnDef",
+    "CostThresholds",
+    "Defaults",
+    "Disclosure",
+    "FunctionPolicy",
+    "LabelDef",
+    "Limits",
+    "Policy",
+    "PolicyError",
+    "Role",
+    "Sensitivity",
+    "TableDef",
+    "TableGrant",
+    "TenantConfig",
+]