bricksentry 0.1.0__py3-none-any.whl

bricksentry/__init__.py

@@ -0,0 +1,26 @@
+"""bricksentry: change-safety for Databricks pipelines and the AI agents that edit them.
+
+One engine, three faces:
+
+* CLI / GitHub Action for data engineers and CI
+* Python API for notebooks, jobs and tests
+* MCP server for AI coding agents
+
+Every face answers the same question: *is this change safe, what will it break,
+and how do I undo it?*
+"""
+
+from bricksentry.assess import Sentry
+from bricksentry.config import Config, TableSpec
+from bricksentry.models import Assessment, ImpactNode, TableDiff, Verdict
+
+__all__ = [
+    "Sentry",
+    "Config",
+    "TableSpec",
+    "Assessment",
+    "ImpactNode",
+    "TableDiff",
+    "Verdict",
+]
+__version__ = "0.1.0"

bricksentry/agent_tools.py

@@ -0,0 +1,139 @@
+"""Agent tools as plain Python functions returning JSON-safe dicts.
+
+The MCP server registers these, but they work in any agent framework
+(LangGraph, OpenAI Agents SDK, Databricks Agent Framework...): import and wrap.
+
+Errors come back as ``{"error": "..."}`` instead of exceptions, so an agent
+can read them and recover.
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import Any
+
+from bricksentry.assess import Sentry
+from bricksentry.explain import explain_diff as _explain_diff
+
+_SENTRY: Sentry | None = None
+
+
+def get_sentry() -> Sentry:
+    global _SENTRY
+    if _SENTRY is None:
+        _SENTRY = Sentry()
+    return _SENTRY
+
+
+def set_sentry(sentry: Sentry | None) -> None:
+    global _SENTRY
+    _SENTRY = sentry
+
+
+def _safe(fn):
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        try:
+            return fn(*args, **kwargs)
+        except Exception as e:
+            return {"error": f"{type(e).__name__}: {e}"}
+
+    return wrapper
+
+
+@_safe
+def whats_downstream(table: str, depth: int = 3) -> dict[str, Any]:
+    """List everything downstream of a table (tables, jobs, notebooks, pipelines,
+    dashboards, queries) with owners and tags, from lineage. Read-only."""
+    s = get_sentry()
+    nodes = s.downstream(table, depth)
+    owner, tags = s.lineage.owner_and_tags(table)
+    return {
+        "table": table,
+        "owner": owner,
+        "tags": tags,
+        "count": len(nodes),
+        "downstream": [n.to_dict() for n in nodes],
+    }
+
+
+@_safe
+def assess_change(
+    tables: list[str],
+    candidates: dict[str, str] | None = None,
+    pr: str | None = None,
+) -> dict[str, Any]:
+    """Assess a change: diff each table's production version against its candidate
+    (built by your change), walk downstream lineage and apply policy.
+    Returns a verdict (pass / warn / needs_human / block), a risk tier, reasons,
+    owners to notify and an assessment id for explain_diff / safe_to_ship. Read-only."""
+    a = get_sentry().assess(list(tables), candidates=candidates, pr=pr)
+    return a.to_dict()
+
+
+@_safe
+def preview_write(sql: str, target: str, keys: list[str] | None = None) -> dict[str, Any]:
+    """Run ONE write statement (INSERT / MERGE / UPDATE / DELETE / CREATE OR REPLACE TABLE)
+    against a disposable sandbox clone of `target`, never production, and report
+    what it would change plus the policy verdict. Use {{target}} in the SQL or the
+    literal table name. The clone is dropped afterwards."""
+    return get_sentry().preview_write(sql, target, keys=keys)
+
+
+@_safe
+def explain_diff(assessment_id: str | None = None, table: str | None = None) -> dict[str, Any]:
+    """Plain-English explanation of what changed, by assessment id or table name. Read-only."""
+    s = get_sentry()
+    if assessment_id:
+        a = s.get_assessment(assessment_id)
+        if a is None:
+            return {"error": f"unknown assessment id {assessment_id}"}
+        return {"assessment_id": a.id, "explanations": {t.table: t.summary for t in a.tables}}
+    if table:
+        d = s.get_diff(table)
+        if d is None:
+            return {"error": f"no diff for {table} yet; call assess_change or preview_write first"}
+        return {"table": table, "explanation": _explain_diff(d)}
+    return {"error": "pass assessment_id or table"}
+
+
+@_safe
+def safe_to_ship(assessment_id: str) -> dict[str, Any]:
+    """Final gate for an assessment: returns safe_to_ship (bool), the verdict,
+    the rules that fired and the next step the agent should take. Read-only."""
+    a = get_sentry().get_assessment(assessment_id)
+    if a is None:
+        return {"error": f"unknown assessment id {assessment_id}"}
+    v = a.verdict.value
+    next_step = {
+        "pass": "Proceed: open or merge the pull request.",
+        "warn": "Proceed, but include the warnings in the pull request description.",
+        "needs_human": "Do not merge. Open the pull request with the report and request review from: "
+        + (", ".join(a.notify) or "a data owner") + ".",
+        "block": "Stop. Fix the blocking issues listed in rule_hits, rebuild the candidate and assess again.",
+    }[v]
+    return {
+        "assessment_id": a.id,
+        "safe_to_ship": v in ("pass", "warn"),
+        "verdict": v,
+        "risk_tier": a.verdict.risk_tier,
+        "rule_hits": [h.to_dict() for h in a.hits],
+        "notify": a.notify,
+        "next_step": next_step,
+    }
+
+
+@_safe
+def rollback_plan(table: str) -> dict[str, Any]:
+    """The exact SQL to restore `table` to its current version if a change goes wrong.
+    Capture it BEFORE applying a change. Read-only."""
+    s = get_sentry()
+    return {
+        "table": table,
+        "current_version": s.backend.current_version(table),
+        "rollback_sql": s.rollback_plan(table),
+    }
+
+
+TOOLS = [whats_downstream, assess_change, preview_write, explain_diff, safe_to_ship, rollback_plan]
+READ_ONLY = {"whats_downstream", "assess_change", "explain_diff", "safe_to_ship", "rollback_plan"}

bricksentry/assess.py

@@ -0,0 +1,165 @@
+"""The ``Sentry``: one object that the CLI, the MCP server and Python users all share."""
+
+from __future__ import annotations
+
+import uuid
+from typing import Any
+
+from bricksentry.backends.base import Backend, make_backend
+from bricksentry.config import Config
+from bricksentry.diff import diff_tables
+from bricksentry.explain import explain_diff, explain_impact
+from bricksentry.lineage import LineageProvider, make_lineage
+from bricksentry.models import Assessment, ImpactNode, TableAssessment, TableDiff, Verdict
+from bricksentry.policy import Policy
+from bricksentry.preview import rewrite_for_preview, sandbox_name
+
+
+class Sentry:
+    def __init__(
+        self,
+        config: Config | None = None,
+        backend: Backend | None = None,
+        lineage: LineageProvider | None = None,
+        policy: Policy | None = None,
+    ):
+        self.config = config or Config.load()
+        self.backend = backend or make_backend(self.config.backend)
+        self.lineage = lineage or make_lineage(self.config.lineage, self.backend)
+        self.policy = policy or Policy(self.config.policy)
+        self._assessments: dict[str, Assessment] = {}
+        self._diffs: dict[str, TableDiff] = {}
+
+    # --- building blocks --------------------------------------------------
+    def diff(
+        self,
+        base: str,
+        candidate: str,
+        keys: list[str] | None = None,
+        table: str | None = None,
+        sample_pct: float | None = None,
+    ) -> TableDiff:
+        spec = self.config.spec(table or base)
+        d = diff_tables(
+            self.backend,
+            base,
+            candidate,
+            keys or spec.keys,
+            table=table or base,
+            sample_pct=sample_pct if sample_pct is not None else self.config.sample_pct,
+            segment_by=spec.segment_by,
+            max_samples=self.config.max_samples,
+        )
+        self._diffs[d.table] = d
+        return d
+
+    def downstream(self, table: str, depth: int | None = None) -> list[ImpactNode]:
+        return self.lineage.downstream(table, depth or int(self.config.lineage.get("depth", 3)))
+
+    def rollback_plan(self, table: str) -> str:
+        v = self.backend.current_version(table)
+        if v is None:
+            return (
+                f"No Delta history available for {table} on this backend; "
+                "take a snapshot (CREATE TABLE ... CLONE) before applying the change."
+            )
+        return f"RESTORE TABLE {self.backend.qualified(table)} TO VERSION AS OF {v}"
+
+    # --- the main entry point ---------------------------------------------
+    def assess(
+        self,
+        tables: list[str],
+        candidates: dict[str, str] | None = None,
+        pr: str | None = None,
+        diffs: dict[str, TableDiff] | None = None,
+    ) -> Assessment:
+        """Assess a change touching ``tables``.
+
+        For each table: diff base vs candidate (unless a diff is supplied),
+        walk downstream lineage, evaluate policy and record a rollback plan.
+        """
+        candidates = candidates or {}
+        diffs = diffs or {}
+        results: list[TableAssessment] = []
+        notify: set[str] = set()
+        for table in tables:
+            diff: TableDiff | None = diffs.get(table)
+            err: str | None = None
+            if diff is None:
+                cand = candidates.get(table) or self.config.candidate_for(table, pr)
+                try:
+                    if not self.backend.table_exists(cand):
+                        raise ValueError(f"candidate table {cand} not found (did the CI build run?)")
+                    diff = self.diff(table, cand, table=table)
+                except Exception as e:  # report, don't crash: an agent needs a verdict
+                    err = str(e)
+            try:
+                nodes = self.downstream(table)
+            except Exception as e:
+                nodes = []
+                err = err or f"lineage lookup failed: {e}"
+            owner, tags = self.lineage.owner_and_tags(table)
+            verdict, hits = self.policy.evaluate(table, diff, nodes, tags, diff_error=err)
+            summary = (explain_diff(diff) if diff else f"{table}: diff unavailable ({err}).") + " " + explain_impact(table, nodes)
+            ta = TableAssessment(
+                table=table,
+                diff=diff,
+                downstream=nodes,
+                owner=owner,
+                tags=tags,
+                rollback=self.rollback_plan(table),
+                verdict=verdict,
+                hits=hits,
+                summary=summary,
+            )
+            results.append(ta)
+            if verdict.rank >= self.policy.notify_on.rank:
+                if owner:
+                    notify.add(owner)
+                for n in nodes:
+                    if n.owner and n.kind == "table":
+                        notify.add(n.owner)
+        a = Assessment(
+            id=uuid.uuid4().hex[:12],
+            tables=results,
+            verdict=Verdict.worst([t.verdict for t in results]),
+            notify=sorted(notify),
+        )
+        self._assessments[a.id] = a
+        return a
+
+    def preview_write(
+        self,
+        sql: str,
+        target: str,
+        keys: list[str] | None = None,
+        keep: bool = False,
+    ) -> dict[str, Any]:
+        """Run a write against a sandbox clone of ``target`` and return the assessment of its effect."""
+        spec = self.config.spec(target)
+        keys = keys or spec.keys
+        if not keys:
+            raise ValueError(f"No keys configured for {target}; pass keys=[...] to preview_write")
+        clone = sandbox_name(self.config.sandbox_schema, target)
+        rewritten = rewrite_for_preview(sql, target, clone, self.backend.quote)
+        schema = ".".join(clone.split(".")[:-1])
+        if schema:
+            self.backend.ensure_schema(schema)
+        self.backend.clone(target, clone)
+        try:
+            self.backend.execute(rewritten)
+            diff = self.diff(target, clone, keys=keys, table=target)
+            assessment = self.assess([target], diffs={target: diff})
+        finally:
+            if not keep:
+                self.backend.drop(clone)
+        out = assessment.to_dict()
+        out["preview"] = {"sandbox_table": clone if keep else None, "executed_sql": rewritten}
+        return out
+
+    # --- lookups for agents -------------------------------------------------
+    def get_assessment(self, assessment_id: str) -> Assessment | None:
+        return self._assessments.get(assessment_id)
+
+    def get_diff(self, table: str) -> TableDiff | None:
+        return self._diffs.get(table)

bricksentry/backends/__init__.py

@@ -0,0 +1,3 @@
+from bricksentry.backends.base import Backend, make_backend
+
+__all__ = ["Backend", "make_backend"]

bricksentry/backends/base.py

@@ -0,0 +1,96 @@
+"""Backend interface. All heavy lifting is pushed down as SQL to the engine
+(Databricks SQL warehouse, or DuckDB locally), so bricksentry itself needs no Spark
+session and stays cheap enough for an agent to call in a loop."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+NUMERIC_PREFIXES = (
+    "tinyint", "smallint", "int", "bigint", "long", "short", "byte", "integer",
+    "float", "double", "real", "decimal", "numeric", "hugeint", "ubigint", "uinteger",
+    "usmallint", "utinyint",
+)
+
+
+class Backend(ABC):
+    name: str = "base"
+    supports_time_travel: bool = False
+
+    # --- execution ---------------------------------------------------------
+    @abstractmethod
+    def query(self, sql: str) -> list[dict[str, Any]]:
+        """Run SQL and return rows as dicts."""
+
+    def execute(self, sql: str) -> None:
+        self.query(sql)
+
+    def scalar(self, sql: str) -> Any:
+        rows = self.query(sql)
+        if not rows:
+            return None
+        return next(iter(rows[0].values()))
+
+    # --- dialect -----------------------------------------------------------
+    @abstractmethod
+    def quote(self, ident: str) -> str: ...
+
+    def qualified(self, table: str) -> str:
+        return ".".join(self.quote(p) for p in table.split("."))
+
+    @property
+    def string_type(self) -> str:
+        return "STRING"
+
+    @abstractmethod
+    def sample_predicate(self, keys: list[str], pct: float, alias: str | None = None) -> str:
+        """Deterministic hash sample on the key columns, identical on both sides."""
+
+    @abstractmethod
+    def days_ago(self, days: int) -> str:
+        """SQL expression for the date ``days`` days ago."""
+
+    # --- metadata ----------------------------------------------------------
+    @abstractmethod
+    def columns(self, table: str) -> dict[str, str]:
+        """Ordered mapping column -> lower-cased type name."""
+
+    def table_exists(self, table: str) -> bool:
+        try:
+            self.columns(table)
+            return True
+        except Exception:
+            return False
+
+    @abstractmethod
+    def clone(self, source: str, target: str) -> None:
+        """Make a cheap copy of ``source`` at ``target`` (Delta SHALLOW CLONE on Databricks)."""
+
+    def drop(self, table: str) -> None:
+        self.execute(f"DROP TABLE IF EXISTS {self.qualified(table)}")
+
+    def ensure_schema(self, schema: str) -> None:
+        self.execute(f"CREATE SCHEMA IF NOT EXISTS {self.qualified(schema)}")
+
+    def current_version(self, table: str) -> int | None:
+        return None
+
+    @staticmethod
+    def is_numeric(type_name: str) -> bool:
+        t = type_name.lower()
+        return t.startswith(NUMERIC_PREFIXES)
+
+
+def make_backend(spec: dict[str, Any] | None) -> Backend:
+    spec = dict(spec or {"type": "duckdb"})
+    kind = spec.pop("type", "duckdb")
+    if kind == "duckdb":
+        from bricksentry.backends.duckdb_backend import DuckDBBackend
+
+        return DuckDBBackend(**spec)
+    if kind == "databricks":
+        from bricksentry.backends.databricks_backend import DatabricksBackend
+
+        return DatabricksBackend(**spec)
+    raise ValueError(f"Unknown backend type: {kind!r} (expected 'databricks' or 'duckdb')")

bricksentry/backends/databricks_backend.py

@@ -0,0 +1,161 @@
+"""Databricks backend: runs SQL on a SQL warehouse via the Statement Execution API.
+
+Auth uses the standard Databricks SDK chain (env vars, ~/.databrickscfg profile,
+OAuth, Azure/GCP identity...). Nothing in bricksentry handles credentials itself.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from typing import Any
+
+from bricksentry.backends.base import Backend
+
+_INT_TYPES = {"BYTE", "SHORT", "INT", "LONG"}
+_FLOAT_TYPES = {"FLOAT", "DOUBLE", "DECIMAL"}
+
+
+class BricksentryQueryError(RuntimeError):
+    pass
+
+
+class DatabricksBackend(Backend):
+    name = "databricks"
+    supports_time_travel = True
+
+    def __init__(
+        self,
+        warehouse_id: str | None = None,
+        profile: str | None = None,
+        client: Any = None,
+        timeout_s: int = 600,
+        **_: Any,
+    ):
+        self.warehouse_id = warehouse_id or os.environ.get("DATABRICKS_WAREHOUSE_ID")
+        if not self.warehouse_id:
+            raise ValueError(
+                "Databricks backend needs a SQL warehouse: set backend.warehouse_id in "
+                "bricksentry.yml or the DATABRICKS_WAREHOUSE_ID env var."
+            )
+        if client is None:
+            try:
+                from databricks.sdk import WorkspaceClient
+            except ImportError as e:  # pragma: no cover
+                raise ImportError("Databricks backend needs `pip install bricksentry[databricks]`") from e
+            client = WorkspaceClient(profile=profile) if profile else WorkspaceClient()
+        self.client = client
+        self.timeout_s = timeout_s
+
+    # --- execution ---------------------------------------------------------
+    def query(self, sql: str) -> list[dict[str, Any]]:
+        api = self.client.statement_execution
+        resp = api.execute_statement(
+            statement=sql,
+            warehouse_id=self.warehouse_id,
+            wait_timeout="30s",
+            **self._format_kwargs(),
+        )
+        deadline = time.monotonic() + self.timeout_s
+        while _state(resp) in ("PENDING", "RUNNING"):
+            if time.monotonic() > deadline:
+                try:
+                    api.cancel_execution(resp.statement_id)
+                finally:
+                    raise BricksentryQueryError(f"Query timed out after {self.timeout_s}s")
+            time.sleep(1.0)
+            resp = api.get_statement(resp.statement_id)
+
+        state = _state(resp)
+        if state != "SUCCEEDED":
+            err = getattr(getattr(resp, "status", None), "error", None)
+            msg = getattr(err, "message", None) or str(err) or state
+            raise BricksentryQueryError(f"Databricks query {state}: {msg}\nSQL: {sql[:500]}")
+
+        manifest = resp.manifest
+        if manifest is None or manifest.schema is None or not manifest.schema.columns:
+            return []
+        cols = manifest.schema.columns
+        names = [c.name for c in cols]
+        types = [_type_name(c) for c in cols]
+
+        rows: list[list[Any]] = []
+        chunk = resp.result
+        while chunk is not None:
+            rows.extend(chunk.data_array or [])
+            nxt = getattr(chunk, "next_chunk_index", None)
+            if nxt is None:
+                break
+            chunk = api.get_statement_result_chunk_n(resp.statement_id, nxt)
+        return [{n: _convert(v, t) for n, v, t in zip(names, row, types)} for row in rows]
+
+    @staticmethod
+    def _format_kwargs() -> dict[str, Any]:
+        try:
+            from databricks.sdk.service.sql import Disposition, Format
+
+            return {"disposition": Disposition.INLINE, "format": Format.JSON_ARRAY}
+        except ImportError:  # pragma: no cover - fake clients in tests
+            return {}
+
+    # --- dialect -----------------------------------------------------------
+    def quote(self, ident: str) -> str:
+        return "`" + ident.replace("`", "``") + "`"
+
+    def sample_predicate(self, keys: list[str], pct: float, alias: str | None = None) -> str:
+        pre = f"{alias}." if alias else ""
+        cols = ", ".join(f"{pre}{self.quote(k)}" for k in keys)
+        return f"(abs(xxhash64({cols})) % 10000) < {int(pct * 100)}"
+
+    def days_ago(self, days: int) -> str:
+        return f"date_sub(current_date(), {int(days)})"
+
+    # --- metadata ----------------------------------------------------------
+    def columns(self, table: str) -> dict[str, str]:
+        out: dict[str, str] = {}
+        for r in self.query(f"DESCRIBE TABLE {self.qualified(table)}"):
+            name = (r.get("col_name") or "").strip()
+            if not name or name.startswith("#"):
+                break  # partition / detail sections follow
+            out[name] = str(r.get("data_type", "")).lower()
+        return out
+
+    def clone(self, source: str, target: str) -> None:
+        self.execute(
+            f"CREATE OR REPLACE TABLE {self.qualified(target)} SHALLOW CLONE {self.qualified(source)}"
+        )
+
+    def current_version(self, table: str) -> int | None:
+        try:
+            rows = self.query(f"DESCRIBE HISTORY {self.qualified(table)} LIMIT 1")
+        except Exception:
+            return None
+        if not rows:
+            return None
+        v = rows[0].get("version")
+        return int(v) if v is not None else None
+
+
+def _state(resp: Any) -> str:
+    st = getattr(getattr(resp, "status", None), "state", None)
+    return getattr(st, "value", st) or "UNKNOWN"
+
+
+def _type_name(col: Any) -> str:
+    t = getattr(col, "type_name", None)
+    return str(getattr(t, "value", t) or "STRING").upper()
+
+
+def _convert(v: Any, t: str) -> Any:
+    if v is None:
+        return None
+    try:
+        if t in _INT_TYPES:
+            return int(v)
+        if t in _FLOAT_TYPES:
+            return float(v)
+        if t == "BOOLEAN":
+            return str(v).lower() == "true"
+    except (TypeError, ValueError):
+        return v
+    return v

bricksentry/backends/duckdb_backend.py

@@ -0,0 +1,52 @@
+"""DuckDB backend: for local development, tests, demos and offline CI."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from bricksentry.backends.base import Backend
+
+
+class DuckDBBackend(Backend):
+    name = "duckdb"
+    supports_time_travel = False
+
+    def __init__(self, database: str = ":memory:", connection: Any = None, **_: Any):
+        if connection is None:
+            try:
+                import duckdb
+            except ImportError as e:  # pragma: no cover
+                raise ImportError("DuckDB backend needs `pip install bricksentry[duckdb]`") from e
+            connection = duckdb.connect(database)
+        self.con = connection
+
+    def query(self, sql: str) -> list[dict[str, Any]]:
+        cur = self.con.execute(sql)
+        if cur.description is None:
+            return []
+        cols = [d[0] for d in cur.description]
+        return [dict(zip(cols, row)) for row in cur.fetchall()]
+
+    def quote(self, ident: str) -> str:
+        return '"' + ident.replace('"', '""') + '"'
+
+    @property
+    def string_type(self) -> str:
+        return "VARCHAR"
+
+    def sample_predicate(self, keys: list[str], pct: float, alias: str | None = None) -> str:
+        pre = f"{alias}." if alias else ""
+        cols = ", ".join(f"{pre}{self.quote(k)}" for k in keys)
+        return f"(hash({cols}) % 10000) < {int(pct * 100)}"
+
+    def days_ago(self, days: int) -> str:
+        return f"(current_date - INTERVAL {int(days)} DAY)"
+
+    def columns(self, table: str) -> dict[str, str]:
+        rows = self.query(f"DESCRIBE {self.qualified(table)}")
+        return {r["column_name"]: str(r["column_type"]).lower() for r in rows}
+
+    def clone(self, source: str, target: str) -> None:
+        self.execute(
+            f"CREATE OR REPLACE TABLE {self.qualified(target)} AS SELECT * FROM {self.qualified(source)}"
+        )

bricksentry/changes.py

@@ -0,0 +1,34 @@
+"""Work out which tables a pull request touches.
+
+Each table in ``bricksentry.yml`` lists the repo paths that produce it
+(notebooks, SQL files, Databricks Asset Bundle resource YAML...). A changed
+file matching a table's globs marks that table as changed.
+"""
+
+from __future__ import annotations
+
+import fnmatch
+import subprocess
+
+from bricksentry.config import Config
+
+
+def changed_files(base_ref: str = "origin/main", cwd: str | None = None) -> list[str]:
+    out = subprocess.run(
+        ["git", "diff", "--name-only", f"{base_ref}...HEAD"],
+        capture_output=True,
+        text=True,
+        cwd=cwd,
+        check=True,
+    ).stdout
+    return [line.strip() for line in out.splitlines() if line.strip()]
+
+
+def tables_for_files(cfg: Config, files: list[str]) -> list[str]:
+    hits: list[str] = []
+    for name, spec in cfg.tables.items():
+        for pat in spec.paths:
+            if any(fnmatch.fnmatch(f, pat) for f in files):
+                hits.append(name)
+                break
+    return hits