vericto 0.2.0__py3-none-any.whl

vericto/__init__.py

@@ -0,0 +1,55 @@
+"""Vericto: check every SQL statement with Vericto before it runs.
+
+The core is driver-agnostic:
+
+* :class:`Vericto` (the guard) applies a verdict: raise on block, run the masked
+  rewrite, warn on flag, and the fail mode when Vericto cannot answer.
+* :class:`Evaluator` is how a verdict is obtained; :class:`RuntimeApiEvaluator`
+  asks the Vericto Runtime API.
+
+Driver wrappers live in submodules so that importing ``vericto`` never imports a
+database driver: :mod:`vericto.snowflake`, :mod:`vericto.sqlite`,
+:mod:`vericto.sqlalchemy` and the generic :mod:`vericto.dbapi`.
+"""
+
+from ._version import __version__
+from .dbapi import wrap
+from .dialects import KNOWN_DIALECTS, normalize_dialect
+from .errors import (
+    VerictoBlocked,
+    VerictoConfigError,
+    VerictoError,
+    VerictoFailOpenWarning,
+    VerictoFlaggedWarning,
+    VerictoQuotaExceeded,
+    VerictoUnavailable,
+    VerictoUnsupportedDialect,
+)
+from .evaluator import Evaluator
+from .guard import Vericto
+from .placeholders import PlaceholderMappingError
+from .runtime_api import DEFAULT_API_URL, RuntimeApiEvaluator
+from .verdict import Decision, Finding, Verdict
+
+__all__ = [
+    "DEFAULT_API_URL",
+    "KNOWN_DIALECTS",
+    "Decision",
+    "Evaluator",
+    "Finding",
+    "PlaceholderMappingError",
+    "RuntimeApiEvaluator",
+    "Verdict",
+    "Vericto",
+    "VerictoBlocked",
+    "VerictoConfigError",
+    "VerictoError",
+    "VerictoFailOpenWarning",
+    "VerictoFlaggedWarning",
+    "VerictoQuotaExceeded",
+    "VerictoUnavailable",
+    "VerictoUnsupportedDialect",
+    "__version__",
+    "normalize_dialect",
+    "wrap",
+]

vericto/__main__.py

@@ -0,0 +1,18 @@
+"""``python -m vericto dbt check ...``: the same command as ``vericto-dbt``."""
+
+from __future__ import annotations
+
+import sys
+
+
+def main(argv: list[str]) -> int:
+    if argv and argv[0] == "dbt":
+        from .dbt import main as dbt_main
+
+        return dbt_main(argv[1:])
+    sys.stderr.write("usage: python -m vericto dbt check [options]  (see: python -m vericto dbt --help)\n")
+    return 2
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

vericto/_agent.py

@@ -0,0 +1,153 @@
+"""Shared pieces of the LLM-framework integrations (LangChain, LlamaIndex).
+
+* :func:`block_message`: the text an agent receives instead of rows when Vericto
+  blocks its query, written so the model can correct itself.
+* :func:`describe_verdict`: a one-paragraph verdict for check tools.
+* :class:`EngineGate`: checks the SQL an integration sends through a SQLAlchemy
+  engine, and only that SQL (not schema reflection, not the framework's own
+  ``SET search_path``), so a query costs one evaluation.
+
+Nothing here imports LangChain or LlamaIndex.
+"""
+
+from __future__ import annotations
+
+import contextvars
+from contextlib import contextmanager
+from typing import Any, Collection, Iterator, Optional
+
+from .errors import VerictoBlocked, VerictoConfigError
+from .guard import Vericto
+from .verdict import Decision, Verdict
+
+ON_BLOCK = ("return", "raise")
+
+
+def check_on_block(value: str) -> str:
+    mode = str(value).strip().lower()
+    if mode not in ON_BLOCK:
+        raise VerictoConfigError(f"on_block must be 'return' or 'raise', got {value!r}.")
+    return mode
+
+
+def _rule_label(verdict: Verdict) -> str:
+    rule = verdict.rule_code or "a Vericto rule"
+    return f"{rule}, severity {verdict.severity}" if verdict.severity else rule
+
+
+def _denied(verdict: Verdict) -> Optional[str]:
+    if not verdict.access_denied:
+        return None
+    parts = []
+    for item in verdict.access_denied:
+        if isinstance(item, dict):
+            name = ".".join(str(item[k]) for k in ("schema", "table", "column") if item.get(k))
+            parts.append(name or str(item))
+        else:
+            parts.append(str(item))
+    return ", ".join(parts)
+
+
+def block_message(blocked: VerictoBlocked) -> str:
+    """What the agent reads when its query is blocked. Never contains row data."""
+    verdict = blocked.verdict
+    lines = [f"Error: Vericto blocked this query ({_rule_label(verdict)}). It was NOT executed."]
+    if verdict.message:
+        lines.append(f"Reason: {verdict.message}")
+    others = [f for f in verdict.findings if f.rule_code != verdict.rule_code]
+    if others:
+        lines.append("Other findings: " + ", ".join(f.rule_code for f in others) + ".")
+    denied = _denied(verdict)
+    if denied:
+        lines.append(f"Not allowed for this agent: {denied}.")
+    if verdict.suggested_fix:
+        lines.append(f"Suggested fix: {verdict.suggested_fix}")
+        lines.append("Rewrite the query following the suggested fix and try again.")
+    else:
+        lines.append("Rewrite the query so that it complies with this rule, or tell the user it cannot be run.")
+    if verdict.event_id:
+        lines.append(f"(Vericto event {verdict.event_id})")
+    return "\n".join(lines)
+
+
+def describe_verdict(verdict: Verdict) -> str:
+    """A readable verdict for a check-only tool."""
+    if verdict.blocked:
+        return block_message(VerictoBlocked(verdict, verdict.sql))
+    if verdict.decision is Decision.UNEVALUATED:
+        return f"UNEVALUATED: Vericto could not evaluate the query ({verdict.message or 'no verdict'})."
+    if verdict.decision is Decision.ALLOWED and not verdict.masked:
+        text = "ALLOWED: Vericto found no problem with this query."
+    elif verdict.decision is Decision.PARSE_ERROR:
+        text = "PARSE_ERROR: Vericto could not parse the query; the workspace policy lets it run."
+        if verdict.parse_error:
+            text += f" Parser: {verdict.parse_error}"
+    else:
+        text = f"{verdict.decision.value}: allowed to run ({_rule_label(verdict)})."
+        if verdict.message:
+            text += f" {verdict.message}"
+    if verdict.masked:
+        text += f" Sensitive columns are masked: the query that will run is: {verdict.rewritten_query}"
+    return text
+
+
+class EngineGate:
+    """A ``before_cursor_execute`` listener that is active only inside :meth:`active`.
+
+    The integrations wrap their query entry points in ``with gate.active():``; the
+    statements SQLAlchemy sends to the driver meanwhile are enforced (blocked ones
+    raise :class:`~vericto.VerictoBlocked`, masked ones run as the rewrite).
+    Statements in ``internal`` (fixed SQL the framework issues itself, such as
+    ``SET search_path``) are let through unevaluated.
+
+    When ``passthrough`` is true the engine is already checked by
+    :func:`vericto.sqlalchemy.instrument`: no listener is added, so a statement is
+    evaluated once, and blocks still raise from that listener.
+    """
+
+    def __init__(
+        self,
+        engine: Any,
+        guard: Vericto,
+        dialect: str,
+        *,
+        internal: Collection[str] = (),
+        passthrough: bool = False,
+    ) -> None:
+        from sqlalchemy import event
+
+        self.guard = guard
+        self.dialect = dialect
+        self.internal = frozenset(internal)
+        self._active: contextvars.ContextVar[bool] = contextvars.ContextVar(f"vericto_gate_{id(self)}", default=False)
+        self._target = getattr(engine, "sync_engine", engine)
+        self._paramstyle: Optional[str] = getattr(self._target.dialect, "paramstyle", None)
+        self._listening = False
+        if not passthrough:
+            event.listen(self._target, "before_cursor_execute", self._before_cursor_execute, retval=True)
+            self._listening = True
+
+    def _before_cursor_execute(
+        self, conn: Any, cursor: Any, statement: str, parameters: Any, context: Any, executemany: bool
+    ) -> tuple[str, Any]:
+        if not self._active.get() or statement in self.internal:
+            return statement, parameters
+        has_params = executemany or bool(parameters)
+        final = self.guard.prepare(statement, self.dialect, paramstyle=self._paramstyle if has_params else None)
+        return final, parameters
+
+    @contextmanager
+    def active(self) -> Iterator[None]:
+        token = self._active.set(True)
+        try:
+            yield
+        finally:
+            self._active.reset(token)
+
+    def remove(self) -> None:
+        """Detach the listener (idempotent)."""
+        from sqlalchemy import event
+
+        if self._listening and event.contains(self._target, "before_cursor_execute", self._before_cursor_execute):
+            event.remove(self._target, "before_cursor_execute", self._before_cursor_execute)
+        self._listening = False

vericto/_version.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

vericto/airflow/__init__.py

@@ -0,0 +1,221 @@
+"""Vericto for Apache Airflow (2.9+ and 3.x), on the common-sql provider.
+
+Two ways in, sharing the same check:
+
+* :class:`VerictoSQLExecuteQueryOperator`: a drop-in ``SQLExecuteQueryOperator``
+  that checks every statement before any of them runs::
+
+      from vericto.airflow import VerictoSQLExecuteQueryOperator
+
+      VerictoSQLExecuteQueryOperator(task_id="cleanup", conn_id="warehouse", sql="sql/cleanup.sql")
+
+* :func:`vericto.airflow.policy.guard_sql_tasks`: a cluster policy for
+  ``airflow_local_settings.task_policy`` that guards existing SQL operators
+  without DAG changes.
+
+The SQL checked is the **rendered** SQL (templates are rendered before the check
+runs). A script is split with the hook's own ``split_sql_string`` and every
+statement is evaluated before the first one runs:
+
+* ``BLOCKED``: the task fails with :class:`VerictoTaskBlocked` (an
+  ``AirflowFailException``, so it is not retried: the same SQL would be blocked
+  again) and nothing runs.
+* masked columns: the masked rewrite runs instead.
+* ``FLAGGED`` / ``MONITORED``: runs; logged on the ``vericto`` logger and passed to
+  the guard's ``on_verdict``.
+* Vericto unavailable: ``fail_mode="closed"`` (default) fails the task with
+  :class:`~vericto.VerictoUnavailable`, which is retried like any other error.
+
+The dialect comes from the connection type (``postgres``, ``redshift`` (as
+postgres), ``mysql``, ``mssql``, ``oracle``, ``snowflake``, ``sqlite``), or from
+``vericto_dialect`` in the connection's extra, or from the operator argument.
+
+Requires the ``airflow`` extra: ``pip install "vericto[airflow]"``. The guard is
+built from the ``VERICTO_*`` environment variables of the worker the first time a
+task needs it (DAG parsing never calls Vericto).
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any, List, Mapping, Optional
+
+from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
+
+from ..dialects import normalize_dialect
+from ..errors import VerictoBlocked, VerictoConfigError
+from ..guard import Vericto
+from ..verdict import Verdict
+
+try:  # Airflow 3
+    from airflow.sdk.exceptions import AirflowFailException
+except ImportError:  # Airflow 2
+    from airflow.exceptions import AirflowFailException  # type: ignore[no-redef,attr-defined,unused-ignore]
+
+logger = logging.getLogger("vericto")
+
+#: Airflow connection type -> Vericto dialect.
+CONN_TYPE_DIALECTS: Mapping[str, str] = {
+    "postgres": "postgres",
+    "postgresql": "postgres",
+    "redshift": "postgres",
+    "mysql": "mysql",
+    "mariadb": "mysql",
+    "mssql": "mssql",
+    "oracle": "oracle",
+    "snowflake": "snowflake",
+    "sqlite": "sqlite",
+}
+ON_UNSUPPORTED = ("fail", "skip")
+EXTRA_KEY = "vericto_dialect"
+
+
+class VerictoTaskBlocked(VerictoBlocked, AirflowFailException):
+    """Vericto blocked a statement of the task; nothing ran. Not retried."""
+
+
+_shared: Optional[Vericto] = None
+_shared_lock = threading.Lock()
+
+
+def shared_guard() -> Vericto:
+    """The process-wide guard, built from ``VERICTO_*`` on first use."""
+    global _shared
+    with _shared_lock:
+        if _shared is None:
+            _shared = Vericto()
+        return _shared
+
+
+def _check_on_unsupported(value: str) -> str:
+    if value not in ON_UNSUPPORTED:
+        raise VerictoConfigError(f"on_unsupported must be 'fail' or 'skip', got {value!r}.")
+    return value
+
+
+def dialect_for_hook(hook: Any, override: Optional[str] = None) -> Optional[str]:
+    """The Vericto dialect for a DB hook, or ``None`` when its connection type is not mapped."""
+    if override:
+        return normalize_dialect(override)
+    conn = None
+    try:
+        conn = hook.connection
+    except Exception:  # no connection property on this hook, or it cannot be read
+        logger.debug("vericto: could not read the hook's connection", exc_info=True)
+    if conn is not None:
+        extra = getattr(conn, "extra_dejson", None) or {}
+        if isinstance(extra, Mapping) and extra.get(EXTRA_KEY):
+            return normalize_dialect(str(extra[EXTRA_KEY]))
+    conn_type = (getattr(conn, "conn_type", None) or getattr(hook, "conn_type", None) or "").lower()
+    return CONN_TYPE_DIALECTS.get(conn_type)
+
+
+def _paramstyle(hook: Any, sql: str, parameters: Any) -> Optional[str]:
+    if not parameters:
+        return None
+    if isinstance(parameters, Mapping) and "%(" in sql:
+        return "pyformat"
+    placeholder = getattr(hook, "placeholder", "%s")
+    return "format" if placeholder == "%s" else None
+
+
+def _statements(hook: Any, sql: Any) -> List[str]:
+    if isinstance(sql, str):
+        split = getattr(hook, "split_sql_string", None)
+        parts = split(sql) if callable(split) else [sql]
+        return [p for p in parts if p and p.strip()]
+    if isinstance(sql, (list, tuple)):
+        return [s for s in sql if isinstance(s, str) and s.strip()]
+    raise TypeError(f"vericto: expected the task's sql as str or list of str, got {type(sql).__name__}.")
+
+
+def check_task_sql(
+    task: Any,
+    *,
+    guard: Optional[Vericto] = None,
+    dialect: Optional[str] = None,
+    on_unsupported: str = "fail",
+) -> List[Verdict]:
+    """Check a SQL task's rendered ``sql`` and apply masks to it. Call before it runs.
+
+    Every statement is evaluated before any runs. Raises :class:`VerictoTaskBlocked`
+    on a block, :class:`~vericto.VerictoUnavailable` (fail-closed) or
+    :class:`~vericto.VerictoConfigError`. When a statement is masked, ``task.sql``
+    is replaced with the statements to run.
+    """
+    hook = task.get_db_hook()
+    name = dialect_for_hook(hook, dialect)
+    if name is None:
+        conn_type = getattr(getattr(hook, "connection", None), "conn_type", None) or getattr(hook, "conn_type", "?")
+        message = (
+            f"Vericto has no dialect for the connection type '{conn_type}' of task {task.task_id}. "
+            f"Set '{EXTRA_KEY}' in the connection extra or vericto_dialect= on the operator "
+            f"(one of: postgres, mysql, oracle, mssql, snowflake)."
+        )
+        if on_unsupported == "skip":
+            logger.warning("vericto: %s The task runs unchecked (on_unsupported='skip').", message)
+            return []
+        raise VerictoConfigError(message + " Nothing was executed.")
+
+    g = guard if guard is not None else shared_guard()
+    g.require_dialect(name)
+    sql = task.sql
+    statements = _statements(hook, sql)
+    parameters = getattr(task, "parameters", None)
+    verdicts: List[Verdict] = []
+    for statement in statements:
+        try:
+            verdicts.append(g.enforce(statement, name, paramstyle=_paramstyle(hook, statement, parameters)))
+        except VerictoBlocked as exc:
+            raise VerictoTaskBlocked(exc.verdict, exc.sql) from None
+
+    if any(v.masked for v in verdicts):
+        rewritten = [v.sql_to_execute for v in verdicts]
+        task.sql = rewritten[0] if isinstance(sql, str) and len(rewritten) == 1 else rewritten
+        logger.info("vericto: task %s runs the masked rewrite of %d statement(s)", task.task_id, len(rewritten))
+    return verdicts
+
+
+class VerictoSQLExecuteQueryOperator(SQLExecuteQueryOperator):
+    """``SQLExecuteQueryOperator`` that checks the rendered SQL with Vericto first.
+
+    Args:
+        vericto_guard: the :class:`~vericto.Vericto` to use; the process-wide guard
+            from ``VERICTO_*`` when omitted.
+        vericto_dialect: override the dialect read from the connection.
+        vericto_on_unsupported: ``"fail"`` (default): a connection type Vericto has
+            no dialect for fails the task before anything runs. ``"skip"``: run it
+            unchecked, with a warning.
+
+    Every other argument is ``SQLExecuteQueryOperator``'s.
+    """
+
+    def __init__(
+        self,
+        *,
+        vericto_guard: Optional[Vericto] = None,
+        vericto_dialect: Optional[str] = None,
+        vericto_on_unsupported: str = "fail",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.vericto_guard = vericto_guard
+        self.vericto_dialect = vericto_dialect
+        self.vericto_on_unsupported = _check_on_unsupported(vericto_on_unsupported)
+
+    def execute(self, context: Any) -> Any:
+        check_task_sql(
+            self, guard=self.vericto_guard, dialect=self.vericto_dialect, on_unsupported=self.vericto_on_unsupported
+        )
+        return super().execute(context)  # type: ignore[no-untyped-call,unused-ignore]  # untyped upstream
+
+
+__all__ = [
+    "CONN_TYPE_DIALECTS",
+    "VerictoSQLExecuteQueryOperator",
+    "VerictoTaskBlocked",
+    "check_task_sql",
+    "dialect_for_hook",
+    "shared_guard",
+]

vericto/airflow/policy.py

@@ -0,0 +1,85 @@
+"""A cluster policy that guards existing SQL tasks with Vericto, without DAG changes.
+
+In ``airflow_local_settings.py`` (on ``PYTHONPATH`` or in ``$AIRFLOW_HOME/config``)
+of the DAG processor **and the workers**::
+
+    from vericto.airflow.policy import guard_sql_tasks
+
+    def task_policy(task):
+        guard_sql_tasks(task)
+
+Every common-sql operator (``SQLExecuteQueryOperator`` and its subclasses, the
+``SQL*Check`` operators, provider operators built on them) then checks its
+rendered SQL with Vericto right before it runs, exactly as
+:class:`~vericto.airflow.VerictoSQLExecuteQueryOperator` does. The check is
+installed as the task's ``pre_execute`` hook (an existing hook still runs, first),
+so the operator class, its serialization and the UI are unchanged.
+
+Airflow applies ``task_policy`` whenever a DAG file is parsed, including on the
+worker that runs the task, so the guard is in place where the SQL executes.
+Parsing never calls Vericto; the guard is built from the worker's ``VERICTO_*``
+environment on the first checked task.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Optional, Tuple, Type
+
+from airflow.providers.common.sql.operators.sql import BaseSQLOperator
+
+from ..guard import Vericto
+from . import VerictoSQLExecuteQueryOperator, _check_on_unsupported, check_task_sql
+
+MARKER = "_vericto_guarded"
+
+
+def guard_sql_tasks(
+    task: Any,
+    *,
+    guard: Optional[Vericto] = None,
+    dialect: Optional[str] = None,
+    on_unsupported: str = "fail",
+    operators: Tuple[Type[Any], ...] = (BaseSQLOperator,),
+    when: Optional[Callable[[Any], bool]] = None,
+) -> bool:
+    """Guard ``task`` if it is a SQL operator. Returns whether it is now guarded.
+
+    Args:
+        task: the task ``task_policy`` receives.
+        guard: the :class:`~vericto.Vericto` to use; the process-wide one from
+            ``VERICTO_*`` when omitted.
+        dialect: force a dialect instead of reading it from the connection.
+        on_unsupported: ``"fail"`` (default) fails a task whose connection type has
+            no Vericto dialect, before anything runs; ``"skip"`` runs it unchecked
+            with a warning. Set ``vericto_dialect`` in a connection's extra to map
+            ODBC/JDBC connections.
+        operators: the operator classes to guard (default: every common-sql operator).
+        when: an optional predicate on the task (e.g. by ``conn_id`` or DAG tags).
+
+    Idempotent: a task is guarded once, and
+    :class:`~vericto.airflow.VerictoSQLExecuteQueryOperator` tasks are left alone
+    (they check themselves).
+    """
+    mode = _check_on_unsupported(on_unsupported)
+    if not isinstance(task, operators) or isinstance(task, VerictoSQLExecuteQueryOperator):
+        return False
+    if getattr(task, MARKER, False):
+        return True
+    if when is not None and not when(task):
+        return False
+
+    previous = getattr(task, "_pre_execute_hook", None)
+
+    def vericto_pre_execute(context: Any) -> None:
+        if previous is not None:
+            previous(context)
+        # The task being executed: the rendered copy Airflow made for this run.
+        running = context.get("task", task) if hasattr(context, "get") else task
+        check_task_sql(running, guard=guard, dialect=dialect, on_unsupported=mode)
+
+    task._pre_execute_hook = vericto_pre_execute
+    setattr(task, MARKER, True)
+    return True
+
+
+__all__ = ["guard_sql_tasks"]